util.py

""" Utilities to make gamefixes easier
"""

import configparser
import os
import sys
import re
import shutil
import signal
import tarfile
import zipfile
import subprocess
import urllib.request
import functools
from .logger import log
from .steamhelper import install_app
from . import config

try:
    import __main__ as protonmain
except ImportError:
    log.warn('Unable to hook into Proton main script environment')

# pylint: disable=unreachable

def which(appname):
    """ Returns the full path of an executable in $PATH
    """

    for path in os.environ['PATH'].split(os.pathsep):
        fullpath = os.path.join(path, appname)
        if os.path.exists(fullpath) and os.access(fullpath, os.X_OK):
            return fullpath
    log.warn(str(appname) + 'not found in $PATH')
    return None


def protondir():
    """ Returns the path to proton
    """

    proton_dir = os.path.dirname(sys.argv[0])
    return proton_dir


def protonprefix():
    """ Returns the wineprefix used by proton
    """

    return os.path.join(
        os.environ['STEAM_COMPAT_DATA_PATH'],
        'pfx/')


def protonnameversion():
    """ Returns the version of proton from sys.argv[0]
    """

    version = re.search('Proton ([0-9]*\\.[0-9]*)', sys.argv[0])
    if version:
        return version.group(1)
    log.warn('Proton version not parsed from command line')
    return None


def protontimeversion():
    """ Returns the version timestamp of proton from the `version` file
    """

    fullpath = os.path.join(protondir(), 'version')
    try:
        with open(fullpath, 'r') as version:
            for timestamp in version.readlines():
                return int(timestamp.strip())
    except OSError:
        log.warn('Proton version file not found in: ' + fullpath)
        return 0
    log.warn('Proton version not parsed from file: ' + fullpath)
    return 0


def protonversion(timestamp=False):
    """ Returns the version of proton
    """

    if timestamp:
        return protontimeversion()
    return protonnameversion()

def once(func=None, retry=None):
    """ Decorator to use on functions which should only run once in a prefix.
    Error handling:
    By default, when an exception occurs in the decorated function, the
    function is not run again. To change that behavior, set retry to True.
    In that case, when an exception occurs during the decorated function,
    the function will be run again the next time the game is started, until
    the function is run successfully.
    Implementation:
    Uses a file (one per function) in PROTONPREFIX/drive_c/protonfixes/run/
    to track if a function has already been run in this prefix.
    """
    if func is None:
        return functools.partial(once, retry=retry)
    retry = retry if retry else False

    #pylint: disable=missing-docstring
    def wrapper(*args, **kwargs):
        func_id = func.__module__ + "." + func.__name__
        prefix = protonprefix()
        directory = os.path.join(prefix, "drive_c/protonfixes/run/")
        file = os.path.join(directory, func_id)
        if not os.path.exists(directory):
            os.makedirs(directory)
        if os.path.exists(file):
            return

        exception = None
        try:
            func(*args, **kwargs)
        except Exception as exc: #pylint: disable=broad-except
            if retry:
                raise exc
            exception = exc

        open(file, 'a').close()

        if exception:
            raise exception #pylint: disable=raising-bad-type

        return
    return wrapper


def _killhanging():
    """ Kills processes that hang when installing winetricks
    """

    # avoiding an external library as proc should be available on linux
    log.debug('Killing hanging wine processes')
    pids = [pid for pid in os.listdir('/proc') if pid.isdigit()]
    badexes = ['mscorsvw.exe']
    for pid in pids:
        try:
            with open(os.path.join('/proc', pid, 'cmdline'), 'rb') as proc_cmd:
                cmdline = proc_cmd.read()
                for exe in badexes:
                    if exe in cmdline.decode():
                        os.kill(int(pid), signal.SIGKILL)
        except IOError:
            continue

def _forceinstalled(verb):
    """ Records verb into the winetricks.log.forced file
    """
    forced_log = os.path.join(protonprefix(), 'winetricks.log.forced')
    with open(forced_log, 'a') as forcedlog:
        forcedlog.write(verb + '\n')


def _checkinstalled(verb, logfile='winetricks.log'):
    """ Returns True if the winetricks verb is found in the winetricks log
    """

    if not isinstance(verb, str):
        return False

    winetricks_log = os.path.join(protonprefix(), logfile)

    # Check for 'verb=param' verb types
    if len(verb.split('=')) > 1:
        wt_verb = verb.split('=')[0] + '='
        wt_verb_param = verb.split('=')[1]
        wt_is_set = False
        try:
            with open(winetricks_log, 'r') as tricklog:
                for xline in tricklog.readlines():
                    if re.findall(r'^' + wt_verb, xline.strip()):
                        wt_is_set = bool(xline.strip() == wt_verb + wt_verb_param)
            return wt_is_set
        except OSError:
            return False
    # Check for regular verbs
    try:
        with open(winetricks_log, 'r') as tricklog:
            if verb in reversed([x.strip() for x in tricklog.readlines()]):
                return True
    except OSError:
        return False
    return False


def checkinstalled(verb):
    """ Returns True if the winetricks verb is found in the winetricks log
        or in the 'winetricks.log.forced' file
    """

    log.info('Checking if winetricks ' + verb + ' is installed')
    if _checkinstalled(verb, 'winetricks.log.forced'):
        return True
    return _checkinstalled(verb)


def is_custom_verb(verb):
    """ Returns path to custom winetricks verb, if found
    """

    verb_name = verb + '.verb'
    verb_dir = 'verbs'

    # check local custom verbs
    verbpath = os.path.expanduser('~/.config/protonfixes/localfixes/' + verb_dir)
    if os.path.isfile(os.path.join(verbpath, verb_name)):
        log.debug('Using local custom winetricks verb from: ' + verbpath)
        return os.path.join(verbpath, verb_name)

    # check custom verbs
    verbpath = os.path.join(os.path.dirname(__file__), verb_dir)
    if os.path.isfile(os.path.join(verbpath, verb_name)):
        log.debug('Using custom winetricks verb from: ' + verbpath)
        return os.path.join(verbpath, verb_name)

    return False


def protontricks(verb):
    """ Runs winetricks if available
    """

    if not checkinstalled(verb):
        log.info('Installing winetricks ' + verb)
        env = dict(protonmain.g_session.env)
        env['WINEPREFIX'] = protonprefix()
        env['WINE'] = protonmain.g_proton.wine_bin
        env['WINELOADER'] = protonmain.g_proton.wine_bin
        env['WINESERVER'] = protonmain.g_proton.wineserver_bin
        env['WINETRICKS_LATEST_VERSION_CHECK'] = 'disabled'
        env['LD_PRELOAD'] = ''
        
        winetricks_bin = os.path.abspath(__file__).replace('util.py','winetricks')
        winetricks_cmd = [winetricks_bin, '--unattended'] + verb.split(' ')

        # check is verb a custom winetricks verb
        custom_verb = is_custom_verb(verb)
        if custom_verb:
            winetricks_cmd = [winetricks_bin, '--unattended', custom_verb]

        if winetricks_bin is None:
            log.warn('No winetricks was found in $PATH')

        if winetricks_bin is not None:

            log.debug('Using winetricks command: ' + str(winetricks_cmd))

            # make sure proton waits for winetricks to finish
            for idx, arg in enumerate(sys.argv):
                if 'waitforexitandrun' not in arg:
                    sys.argv[idx] = arg.replace('run', 'waitforexitandrun')
                    log.debug(str(sys.argv))

            log.info('Using winetricks verb ' + verb)
            subprocess.call([env['WINESERVER'], '-w'], env=env)
            process = subprocess.Popen(winetricks_cmd, env=env)
            process.wait()
            _killhanging()

            # Check if verb recorded to winetricks log
            if not checkinstalled(verb):
                log.warn('Not recorded as installed: winetricks ' + verb + ', forcing!')
                _forceinstalled(verb)

            log.info('Winetricks complete')
            return True

    return False

def protontricks_proton_5(verb):
    """ Runs winetricks with Proton 5 which is still useful to install some things like .NET"""
    if checkinstalled(verb):
        log.debug("Skipping {} as it is marked as installed".format(verb))
        return

    prefix_path = protonprefix()
    try:
        log.info("Removing the prefix at {} to recreate it with Proton 5".format(prefix_path))
        shutil.rmtree(prefix_path)
    except FileNotFoundError:
        log.warn('The protonprefix folder was not found')

    log.info('Folder Proton 5.0' + str(os.path.join(os.environ['STEAM_COMPAT_DATA_PATH'],'..','..','common','Proton 5.0')))

    wine_path = os.path.join(os.environ['STEAM_COMPAT_DATA_PATH'],'..','..','common','Proton 5.0','dist','bin','wine')
    # If this is being used to install Dotnet for example and it doesn't exist, failing silently might not be enough
    if not os.path.exists(wine_path):
        message = "Ensure Proton 5.0 is installed. No Proton 5.0 was found at the expected path at {}".format(wine_path)
        try_show_gui_error(message)
        raise Exception(message)

    env = dict(protonmain.g_session.env)
    env['WINEPREFIX'] = prefix_path
    env['WINE'] = wine_path
    env['WINELOADER'] = os.path.join(os.environ['STEAM_COMPAT_DATA_PATH'],'..','..','common','Proton 5.0','dist','bin','wine')
    env['WINESERVER'] = os.path.join(os.environ['STEAM_COMPAT_DATA_PATH'],'..','..','common','Proton 5.0','dist','bin','wineserver')
    env['WINEPATH'] = os.path.join(os.environ['STEAM_COMPAT_DATA_PATH'],'..','..','common','Proton 5.0','dist','bin','wine64')
    env['WINETRICKS_LATEST_VERSION_CHECK'] = 'disabled'
    env['LD_PRELOAD'] = ''
    # The reason for the separate Winetricks is newer Winetricks with Proton 5.0 seems to result
    # in the `winecfg` window popping up and things not getting done.
    winetricks_bin = os.path.abspath(__file__).replace('util.py','winetricks_proton5')
    winetricks_cmd = [winetricks_bin, '--unattended', '--force'] + verb.split(' ')


    process = subprocess.Popen(winetricks_cmd, env=env)
    process.wait()

def regedit_add(folder,name=None,type=None,value=None,arch=None):
    """ Add regedit keys
    """

    env = dict(protonmain.g_session.env)
    env['WINEPREFIX'] = protonprefix()
    env['WINE'] = protonmain.g_proton.wine_bin
    env['WINELOADER'] = protonmain.g_proton.wine_bin
    env['WINESERVER'] = protonmain.g_proton.wineserver_bin
    
    if name is not None and type is not None and value is not None:

        # Flag for if we want to force writing to the 64-bit registry sector
        if arch is not None:
            regedit_cmd = ['wine', 'reg' , 'add', folder, '/f', '/v', name, '/t', type, '/d', value, '/reg:64']
        else:
            regedit_cmd = ['wine', 'reg' , 'add', folder, '/f', '/v', name, '/t', type, '/d', value]

        log.info('Adding key: ' + folder)

    else:

        # Flag for if we want to force writing to the 64-bit registry sector
        # We use name here because without the other flags we can't use the arch flag
        if name is not None:
            regedit_cmd = ['wine', 'reg' , 'add', folder, '/f', '/reg:64']
        else:
            regedit_cmd = ['wine', 'reg' , 'add', folder, '/f']

        log.info('Adding key: ' + folder)

    process = subprocess.Popen(regedit_cmd, env=env)
    process.wait()

def replace_command(orig_str, repl_str):
    """ Make a commandline replacement in sys.argv
    """

    log.info('Changing ' + orig_str + ' to ' + repl_str)
    for idx, arg in enumerate(sys.argv):
        if orig_str in arg:
            sys.argv[idx] = arg.replace(orig_str, repl_str)

def append_argument(argument):
    """ Append an argument to sys.argv
    """

    log.info('Adding argument ' + argument)
    sys.argv.append(argument)
    log.debug('New commandline: ' + str(sys.argv))

def set_environment(envvar, value):
    """ Add or override an environment value
    """

    log.info('Adding env: ' + envvar + '=' + value)
    os.environ[envvar] = value
    protonmain.g_session.env[envvar] = value

def del_environment(envvar):
    """ Remove an environment variable
    """

    log.info('Removing env: ' + envvar)
    if envvar in os.environ:
        del os.environ[envvar]
    if envvar in protonmain.g_session.env:
        del protonmain.g_session.env[envvar]

def get_game_install_path():
    """ Game installation path
    """

    log.debug('Detected path to game: ' + os.environ['PWD'])
    # only for `waitforexitandrun` command
    return os.environ['PWD']

def winedll_override(dll, dtype):
    """ Add WINE dll override
    """

    log.info('Overriding ' + dll + '.dll = ' + dtype)
    protonmain.g_session.dlloverrides[dll] = dtype

def disable_nvapi():
    """ Disable WINE nv* dlls
    """

    log.info('Disabling NvAPI')
    winedll_override('nvapi', '')
    winedll_override('nvapi64', '')
    winedll_override('nvcuda', '')
    winedll_override('nvcuvid', '')
    winedll_override('nvencodeapi', '')
    winedll_override('nvencodeapi64', '')

def disable_dxvk():  # pylint: disable=missing-docstring
    set_environment('PROTON_USE_WINED3D', '1')

def disable_esync():  # pylint: disable=missing-docstring
    set_environment('PROTON_NO_ESYNC', '1')

def disable_fsync(): # pylint: disable=missing-docstring
    set_environment('PROTON_NO_FSYNC', '1')

def disable_protonaudioconverter(): # pylint: disable=missing-docstring
    set_environment('GST_PLUGIN_FEATURE_RANK', 'protonaudioconverterbin:NONE')

def use_seccomp(): # pylint: disable=missing-docstring
    set_environment('PROTON_USE_SECCOMP', '1')

@once
def disable_uplay_overlay():
    """Disables the UPlay in-game overlay.
    Creates or appends the UPlay settings.yml file
    with the correct setting to disable the overlay.
    UPlay will overwrite settings.yml on launch, but keep
    this setting.
    """
    config_dir = os.path.join(
        protonprefix(),
        'drive_c/users/steamuser/Local Settings/Application Data/Ubisoft Game Launcher/'
    )

    if not os.path.exists(config_dir):
        os.makedirs(config_dir)

    config_file = os.path.join(config_dir, 'settings.yml')

    if not os.path.isdir(config_dir):
        log.warn(
            'Could not disable UPlay overlay: "'
            + config_dir
            + '" does not exist or is not a directory.'
        )
        return

    if not os.path.isfile(config_file):
        f = open(config_file,"w+")
        f.write("\noverlay:\n  enabled: false\n  forceunhookgame: false\n  fps_enabled: false\n  warning_enabled: false\n")
        f.close
        log.info('Disabled UPlay overlay')
    else:
        try:
            with open(config_file, 'a+') as file:
                file.write("\noverlay:\n  enabled: false\n  forceunhookgame: false\n  fps_enabled: false\n  warning_enabled: false\n")
            log.info('Disabled UPlay overlay')
            return
        except OSError as err:
            log.warn('Could not disable UPlay overlay: ' + err.strerror)

def create_dosbox_conf(conf_file, conf_dict):
    """Create DOSBox configuration file.

    DOSBox accepts multiple configuration files passed with -conf
    option;, each subsequent one overwrites settings defined in
    previous files.
    """
    if os.access(conf_file, os.F_OK):
        return
    conf = configparser.ConfigParser()
    conf.read_dict(conf_dict)
    with open(conf_file, 'w') as file:
        conf.write(file)


def _get_config_full_path(cfile, base_path):
    """ Find game's config file
    """

    # Start from 'user'/'game' directories or absolute path
    if base_path == 'user':
        cfg_path = os.path.join(protonprefix(), 'drive_c/users/steamuser/My Documents', cfile)
    else:
        if base_path == 'game':
            cfg_path = os.path.join(get_game_install_path(), cfile)
        else:
            cfg_path = cfile

    if os.path.exists(cfg_path) and os.access(cfg_path, os.F_OK):
        log.debug('Found config file: ' + cfg_path)
        return cfg_path

    log.warn('Config file not found: ' + cfg_path)
    return False

def create_backup_config(cfg_path):
    """ Create backup config file
    """

    # Backup
    if not os.path.exists(cfg_path + '.protonfixes.bak'):
        log.info('Creating backup for config file')
        shutil.copyfile(cfg_path, cfg_path + '.protonfixes.bak')

def set_ini_options(ini_opts, cfile, encoding, base_path='user'):
    """ Edit game's INI config file
    """
    cfg_path = _get_config_full_path(cfile, base_path)
    if not cfg_path:
        return False

    create_backup_config(cfg_path)

    # set options
    conf = configparser.ConfigParser(empty_lines_in_values=True, allow_no_value=True, strict=False)
    conf.optionxform = str

    conf.read(cfg_path,encoding)

    log.info('Addinging INI options into '+cfile+':\n'+ str(ini_opts))
    conf.read_string(ini_opts)

    with open(cfg_path, 'w') as configfile:
        conf.write(configfile)
    return True

def set_xml_options(base_attibutte, xml_line, cfile, base_path='user'):
    """ Edit game's XML config file
    """
    xml_path = _get_config_full_path(cfile, base_path)
    if not xml_path:
        return False

    create_backup_config(xml_path)

    # set options

    base_size = os.path.getsize(xml_path)
    backup_size = os.path.getsize(xml_path + '.protonfixes.bak')

    if base_size == backup_size:
        ConfigFile = open(xml_path, 'r')
        contents = ConfigFile.readlines()
        LINENUM=0
        for line in contents:
            LINENUM+=1
            if base_attibutte in line:
                log.info('Addinging XML options into '+cfile+':\n'+ str(xml_line))
                contents.insert(LINENUM, xml_line + "\n")
        ConfigFile.close()
        ConfigFile = open(xml_path, 'w')
        for eachitem in contents:
            ConfigFile.write(eachitem)
        ConfigFile.close()
        log.info("Config Patch Applied! \n")

def get_resolution():
    """ Returns screen res width, height
    """

    with open('/sys/class/graphics/fb0/virtual_size', 'r') as res:
        screenx, screeny = map(int, res.read().strip('\n').split(','))

        return(screenx,screeny)

def read_dxvk_conf(cfp):
    """ Add fake [DEFAULT] section to dxvk.conf
    """
    yield '['+ configparser.ConfigParser().default_section +']'
    yield from cfp


def set_dxvk_option(opt, val, cfile='/tmp/protonfixes_dxvk.conf'):
    """ Create custom DXVK config file

    See https://github.com/doitsujin/dxvk/wiki/Configuration for details
    """
    conf = configparser.ConfigParser()
    conf.optionxform = str
    section = conf.default_section
    dxvk_conf = os.path.join(get_game_install_path(), 'dxvk.conf')

    conf.read(cfile)

    if not conf.has_option(section, 'session') or conf.getint(section, 'session') != os.getpid():
        log.info('Creating new DXVK config')
        set_environment('DXVK_CONFIG_FILE', cfile)

        conf = configparser.ConfigParser()
        conf.optionxform = str
        conf.set(section, 'session', str(os.getpid()))

        if os.access(dxvk_conf, os.F_OK):
            conf.read_file(read_dxvk_conf(open(dxvk_conf)))
        log.debug(conf.items(section))

    # set option
    log.info('Addinging DXVK option: '+ str(opt) + ' = ' + str(val))
    conf.set(section, opt, str(val))

    with open(cfile, 'w') as configfile:
        conf.write(configfile)

def install_eac_runtime():
    """ Install Proton Easyanticheat Runtime
    """
    install_app(1826330)

def install_battleye_runtime():
    """ Install Proton BattlEye Runtime
    """
    install_app(1161040)

def install_all_from_tgz(url, path=os.getcwd()):
    """ Install all files from a downloaded tar.gz
    """
 
    cache_dir = config.cache_dir
    tgz_file_name = os.path.basename(url)
    tgz_file_path = os.path.join(cache_dir, tgz_file_name)

    if tgz_file_name not in os.listdir(cache_dir):
        log.info('Downloading ' + tgz_file_name)
        urllib.request.urlretrieve(url, tgz_file_path)
    
    with tarfile.open(tgz_file_path, 'r:gz') as tgz_obj:
        log.info('Extracting ' + tgz_file_name + ' to ' + path)
        tgz_obj.extractall(path)
    
def install_from_zip(url, filename, path=os.getcwd()):
    """ Install a file from a downloaded zip
    """

    if filename in os.listdir(path):
        log.info('File ' + filename + ' found in ' + path)
        return

    cache_dir = config.cache_dir
    zip_file_name = os.path.basename(url)
    zip_file_path = os.path.join(cache_dir, zip_file_name)

    if zip_file_name not in os.listdir(cache_dir):
        log.info('Downloading ' + filename + ' to ' + zip_file_path)
        urllib.request.urlretrieve(url, zip_file_path)

    with zipfile.ZipFile(zip_file_path, 'r') as zip_obj:
        log.info('Extracting ' + filename + ' to ' + path)
        zip_obj.extract(filename, path=path)

def try_show_gui_error(text):
    try:  # in case in-use Python doesn't have tkinter, which is likely
        from tkinter import messagebox
        messagebox.showerror("Proton Fixes", text)
    except Exception as e:
        try:
            subprocess.run(["notify-send", "protonfixes", text])
        except:
            log.info("Failed to show error message with the following text: {}".format(text))

def is_smt_enabled() -> bool:
    """ Returns whether SMT is enabled.
        If the check has failed, False is returned.
    """
    try:
        with open('/sys/devices/system/cpu/smt/active') as smt_file:
            return smt_file.read().strip() == "1"
    except PermissionError:
        log.warn('No permission to read SMT status')
    except OSError as e:
        log.warn(f'SMT status not supported by the kernel (errno: {e.errno})')
    return False

def get_cpu_count() -> int:
    """ Returns the cpu core count, provided by the OS.
        If the request failed, 0 is returned.
    """
    cpu_cores = os.cpu_count()
    if not cpu_cores or cpu_cores <= 0:
        log.warn('Can not read count of logical cpu cores')
        return 0
    return cpu_cores

def set_cpu_topology(core_count: int, ignore_user_setting: bool = False) -> bool:
    """ This sets the cpu topology to a fixed core count.
        By default, a user provided topology is prioritized.
        You can override this behavior by setting `ignore_user_setting`.
    """

    # Don't override the user's settings (except, if we override it)
    user_topo = os.getenv('WINE_CPU_TOPOLOGY')
    if user_topo and not ignore_user_setting:
        log.info(f'Using WINE_CPU_TOPOLOGY set by the user: {user_topo}')
        return False

    # Sanity check
    if not core_count or core_count <= 0:
        log.warn('Only positive core_counts can be used to set cpu topology')
        return False

    # Format (example, 4 cores): 4:0,1,2,3
    cpu_topology = f'{core_count}:{",".join(map(str, range(core_count)))}'
    set_environment('WINE_CPU_TOPOLOGY', cpu_topology)
    log.info(f'Using WINE_CPU_TOPOLOGY: {cpu_topology}')
    return True


def set_cpu_topology_nosmt(core_limit: int = 0, ignore_user_setting: bool = False, threads_per_core: int = 2) -> bool:
    """ This sets the cpu topology to the count of physical cores.
        If SMT is enabled, eg. a 4c8t cpu is limited to 4 logical cores.
        You can limit the core count to the `core_limit` argument.
    """

    # Check first, if SMT is enabled
    if is_smt_enabled() is False:
        log.info('SMT is not active, skipping fix')
        return False

    # Currently (2024) SMT allows 2 threads per core, this might change in the future
    cpu_cores = get_cpu_count() // threads_per_core             # Apply divider
    cpu_cores = max(cpu_cores, min(cpu_cores, core_limit))      # Apply limit
    return set_cpu_topology(cpu_cores, ignore_user_setting)

def set_cpu_topology_limit(core_limit: int, ignore_user_setting: bool = False) -> bool:
    """ This sets the cpu topology to a limited number of logical cores.
        A limit that exceeds the available cores, will be ignored.
    """

    cpu_cores = get_cpu_count()
    if core_limit >= cpu_cores:
        log.info(f'The count of logical cores ({cpu_cores}) is lower than '
                 f'or equal to the set limit ({core_limit}), skipping fix')
        return False

    # Apply the limit
    return set_cpu_topology(core_limit, ignore_user_setting)