scriptutils.py

"""This module has functions for looking into scripts to decide how to launch.

Currently, this is primarily shebang lines. This support is intended to allow
scripts to be somewhat portable between POSIX (where they are natively handled)
and Windows, when launching in Python. They are not intended to provide generic
shebang support, although for historical/compatibility reasons it is possible.

Shebang commands shaped like '/usr/bin/<command>' or '/usr/local/bin/<command>'
will have the command matched to an alias or executable name for detected
runtimes, with the first match being selected.
A command of 'py', 'pyw', 'python' or 'pythonw' will match the default runtime.
If the install manager has been launched in windowed mode, and the selected
alias is not marked as windowed, then the first windowed 'run-for' target will
be substituted (if present - otherwise, it will just not run windowed). Aliases
that map to windowed targets are launched windowed.
If no matching command is found, the default install will be used.

Shebang commands shaped like '/usr/bin/env <command>' will do the same lookup as
above. If no matching command is found, the current PATH environment variable
will be searched for a matching command. It will be launched with a warning,
configuration permitting.

Other shebangs will be treated directly as the command, doing the same lookup
and the same PATH search.

It is not yet implemented, but this is also where a search for PEP 723 inline
script metadata would go. Find the comment mentioning PEP 723 below.
"""

import re

from .logging import LOGGER
from .pathutils import Path, PurePath


class NewEncoding(Exception):
    pass


class NoShebang(Exception):
    pass


def _find_shebang_command(cmd, full_cmd, *, windowed=None):
    sh_cmd = PurePath(full_cmd)
    # HACK: Assuming alias/executable suffix is '.exe' here
    # (But correctly assuming we can't use with_suffix() or .stem)
    if not sh_cmd.match("*.exe"):
        sh_cmd = sh_cmd.with_name(sh_cmd.name + ".exe")

    is_wdefault = sh_cmd.match("pythonw.exe") or sh_cmd.match("pyw.exe")
    is_default = is_wdefault or sh_cmd.match("python.exe") or sh_cmd.match("py.exe")

    # Internal logic error, but non-fatal, if it has no value
    assert windowed is not None

    # Ensure we use the default install for a default name. Otherwise, a
    # "higher" runtime may claim it via an alias, which is not the intent.
    if is_default:
        for i in cmd.get_installs():
            if i.get("default"):
                exe = i["executable"]
                if is_wdefault or windowed:
                    target = [t for t in i.get("run-for", []) if t.get("windowed")]
                    if target:
                        exe = target[0]["target"]
                return {**i, "executable": i["prefix"] / exe}

    for i in cmd.get_installs():
        for a in i.get("alias", ()):
            if sh_cmd.match(a["name"]):
                exe = a["target"]
                LOGGER.debug("Matched alias %s in %s", a["name"], i["id"])
                if windowed and not a.get("windowed"):
                    target = [t for t in i.get("run-for", []) if t.get("windowed")]
                    if target:
                        exe = target[0]["target"]
                        LOGGER.debug("Substituting target %s for windowed=1", exe)
                return {**i, "executable": i["prefix"] / exe}
        if sh_cmd.full_match(PurePath(i["executable"]).name):
            LOGGER.debug("Matched executable name %s in %s", i["executable"], i["id"])
            return i
        if sh_cmd.match(i["executable"]):
            LOGGER.debug("Matched executable %s in %s", i["executable"], i["id"])
            return i

    # Fallback search for 'python[w]<TAG>.exe' shebangs
    if sh_cmd.match("pythonw*.exe"):
        tag = sh_cmd.name[7:-4]
        return cmd.get_install_to_run(f"PythonCore/{tag}", windowed=True)
    if sh_cmd.match("python*.exe"):
        tag = sh_cmd.name[6:-4]
        return cmd.get_install_to_run(f"PythonCore/{tag}")

    raise LookupError


def _find_on_path(cmd, full_cmd):
    import os
    import shutil
    # Recursion prevention
    if os.getenv("__PYTHON_MANAGER_SUPPRESS_ARBITRARY_SHEBANG"):
        raise LookupError
    os.environ["__PYTHON_MANAGER_SUPPRESS_ARBITRARY_SHEBANG"] = "1"

    exe = shutil.which(full_cmd)
    if not exe:
        raise LookupError
    return {
        "display-name": "Shebang command",
        "sort-version": "0.0",
        "executable": Path(exe),
    }


def _parse_shebang(cmd, line, *, windowed=None):
    # For /usr[/local]/bin, we look for a matching alias name.
    shebang = re.match(r"#!\s*/usr/(?:local/)?bin/(?!env\b)([^\\/\s]+).*", line)
    if shebang:
        # Handle the /usr[/local]/bin/python cases
        full_cmd = shebang.group(1)
        LOGGER.debug("Matching shebang: %s", full_cmd)
        try:
            return _find_shebang_command(cmd, full_cmd, windowed=windowed)
        except LookupError:
            LOGGER.warn("A shebang '%s' was found, but could not be matched "
                        "to an installed runtime.", full_cmd)
            LOGGER.warn('If the script does not behave properly, try '
                        'installing the correct runtime with "py install".')
            raise

    # For /usr/bin/env, we look for a matching alias, followed by PATH search.
    # We warn about the PATH search, because we don't know we'll be launching
    # Python at all in this case.
    shebang = re.match(r"#!\s*/usr/bin/env\s+(?:-S\s+)?([^\\/\s]+).*", line)
    if shebang:
        # First do regular install lookup for /usr/bin/env shebangs
        full_cmd = shebang.group(1)
        try:
            return _find_shebang_command(cmd, full_cmd, windowed=windowed)
        except LookupError:
            pass
        # If not, warn and do regular PATH search
        if cmd.shebang_can_run_anything or cmd.shebang_can_run_anything_silently:
            i = _find_on_path(cmd, full_cmd)
            if not cmd.shebang_can_run_anything_silently:
                LOGGER.warn("A shebang '%s' was found but could not be matched "
                            "to an installed runtime, so it will be treated as "
                            "an arbitrary command.", full_cmd)
                LOGGER.warn("To prevent execution of programs that are not "
                            "Python runtimes, set 'shebang_can_run_anything' to "
                            "'false' in your configuration file.")
            return i

        else:
            LOGGER.warn("A shebang '%s' was found, but could not be matched "
                        "to an installed runtime.", full_cmd)
            LOGGER.warn("Arbitrary command execution is disabled. Configure "
                        "'shebang_can_run_anything' to 'true' in your "
                        "configuration file to enable it. "
                        "Launching with default runtime.")
            raise LookupError

    # All other shebangs get treated as arbitrary commands. We warn about
    # this case, because we don't know we'll be launching Python at all.
    shebang = re.match(r"#!\s*(.+)\S*$", line)
    if shebang:
        full_cmd = shebang.group(1)
        # A regular lookup will handle the case where the entire shebang is
        # a valid alias.
        try:
            return _find_shebang_command(cmd, full_cmd, windowed=windowed)
        except LookupError:
            pass
        if cmd.shebang_can_run_anything or cmd.shebang_can_run_anything_silently:
            if not cmd.shebang_can_run_anything_silently:
                LOGGER.warn("A shebang '%s' was found but does not match any "
                            "supported template (e.g. '/usr/bin/python'), so it "
                            "will be treated as an arbitrary command.", full_cmd)
                LOGGER.warn("To prevent execution of programs that are not "
                            "Python runtimes, set 'shebang_can_run_anything' to "
                            "'false' in your configuration file.")
            return _find_on_path(cmd, full_cmd)
        else:
            LOGGER.warn("A shebang '%s' was found, but could not be matched "
                        "to an installed runtime.", full_cmd)
            LOGGER.warn("Arbitrary command execution is disabled. Change "
                        "'shebang_can_run_anything' to 'true' in your "
                        "configuration file to enable it. "
                        "Launching with default runtime.")
            raise LookupError

    raise NoShebang


def _read_script(cmd, script, encoding, *, windowed=None):
    try:
        f = open(script, "r", encoding=encoding, errors="replace")
    except OSError as ex:
        raise LookupError(script) from ex
    with f:
        first_line = next(f, "").rstrip()
        if first_line.startswith("#!"):
            try:
                return _parse_shebang(cmd, first_line, windowed=windowed)
            except LookupError:
                raise LookupError(script) from None
            except NoShebang:
                pass

        coding = re.match(r"\s*#.*coding[=:]\s*([-\w.]+)", first_line)
        if coding and coding.group(1) != encoding:
            raise NewEncoding(coding.group(1))

        # TODO: Parse inline script metadata (PEP 723)
        # This involves finding '# /// script' followed by
        # a line with '# requires-python = <spec>'.
        # That spec needs to be processed as a version constraint, which
        # cmd.get_install_to_run() can handle.
    raise LookupError(script)


def find_install_from_script(cmd, script, *, windowed=False):
    try:
        return _read_script(cmd, script, "utf-8-sig", windowed=windowed)
    except NewEncoding as ex:
        encoding = ex.args[0]
    return _read_script(cmd, script, encoding, windowed=windowed)


def _maybe_quote(a):
    if a[:1] == a[-1:] == '"':
        a = a[1:-1]
    if " " not in a and '"' not in a:
        return a
    if a.endswith("\\"):
        c = len(a) - len(a.rstrip("\\"))
        a += "\\" * c
    if '"' in a:
        bits = []
        for b in a.split('"'):
            if bits:
                bits.append('\\"')
            bits.append(b)
            if b[-1:] == "\\":
                bits.append("\\" * (len(b) - len(b.rstrip("\\"))))
        print(a.split('"'), bits)
        a = ''.join(bits)
    return f'"{a}"' if ' ' in a else a


def quote_args(args):
    """Quotes the provided sequence of arguments preserving all characters.

All backslashes and quotes in the existing arguments will be preserved and will
round-trip through CreateProcess to another Python instance (or another app
using the same parsing rules).

When an argument already starts and ends with a double quote ('"'), they will be
removed and only replaced if necessary.
"""
    return " ".join(_maybe_quote(a) for a in args)


def split_args(arg_string, argv0=False):
    """Splits a single argument string into separate unquoted items.

If argv0 is True, the first argument is parsed as if it is the executable name.
"""
    args = []
    if argv0 and arg_string[:1] == '"':
        a, _, arg_string = arg_string[1:].partition('"')
        args.append(a)

    arg_buffer = None
    quoted_arg = []
    bits = arg_string.strip().split(' ')
    while bits:
        a = bits.pop(0)
        pre, quot, post = a.partition('"')
        if arg_buffer:
            pre = arg_buffer + pre
            arg_buffer = None
        if not quot:
            if quoted_arg:
                quoted_arg.append(pre)
                quoted_arg.append(' ')
            else:
                args.append(pre.replace('\\\\', '\\'))
            continue
        if pre[-1:] == '\\' and (len(pre) - len(pre.rstrip('\\'))) % 2 == 1:
            arg_buffer = pre[:-1] + quot
            if post:
                bits.insert(0, post)
            continue
        elif quoted_arg:
            quoted_arg.append(pre)
            args.append(''.join(quoted_arg).replace('\\\\', '\\'))
            quoted_arg.clear()
            continue

        quoted_arg.append(pre)
        if pre:
            quoted_arg.append(quot)
        if post:
            bits.insert(0, post)
    if quoted_arg:
        args.append(''.join(quoted_arg))
    return args