# Copyright 2021 Offensive Security
# SPDX-license-identifier: GPL-3.0-only

import errno
import filecmp
import getpass
import logging
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path

logger = logging.getLogger(__name__)


def bold(text):
    bold = "\033[1m"
    reset = "\033[0m"
    return bold + text + reset


def say(text):
    print(bold(">>> " + text))
    logger.debug(text)


def say_install_packages(pkgs):
    if isinstance(pkgs, str):
        pkgs = [pkgs]
    if len(pkgs) == 0:
        text = "No package to install"
    elif len(pkgs) == 1:
        text = "Installing package: " + pkgs[0]
    else:
        text = "Installing packages: " + " ".join(pkgs)
    say(text)


def say_remove_packages(pkgs):
    if isinstance(pkgs, str):
        pkgs = [pkgs]
    if len(pkgs) == 0:
        text = "No package to remove"
    elif len(pkgs) == 1:
        text = "Removing package: " + pkgs[0]
    else:
        text = "Removing packages: " + " ".join(pkgs)
    say(text)


def say_install_program(program):
    say(f"Installing program: {program}")


def say_enable_service(service):
    say(f"Enabling service: {service}")


def say_disable_service(service):
    say(f"Disabling service: {service}")


def say_change_login_shell(shell):
    say(f"Setting login shell to '{shell}'")


def backup_filename(filename):
    """Return a backup filename"""
    return filename + ".kali-tweaks-orig"


def tmp_filename(filename):
    """Return a temporary filename"""
    # XXX Use proper tmp function
    return filename + ".tmp"


def _get_kali_tweaks_file(directory, filename):
    """Return the path to a kali-tweaks file"""
    if "RUNNING_FROM_SOURCE_TREE" in os.environ:
        f = os.path.join(os.getcwd(), directory, filename)
        if os.path.isfile(f):
            return f
    else:
        # path is hardcoded in setup.cfg, so it's hardcoded here as well
        f = os.path.join("/usr/lib/kali_tweaks", directory, filename)
        if os.path.isfile(f):
            return f
    raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), f)


def get_datafile(filename):
    """Return the path to a data file"""
    return _get_kali_tweaks_file("data", filename)


def get_helper(filename):
    """Return the path to a helper"""
    return _get_kali_tweaks_file("helpers", filename)


def _run(cmd, interactive=False, log=None):
    """If interactive is set, don't capture the output"""
    if log is None:
        log = True if interactive is True else False
    if log is True:
        print("> $ " + cmd)
    logger.debug("Running cmd: %s", cmd)
    res = subprocess.run(cmd, capture_output=(not interactive), shell=True, text=True)
    logger.debug("Return code: %d", res.returncode)
    # logger.debug("Stdout     :\n%s", res.stdout)
    # logger.debug("Stderr     :\n%s", res.stderr)
    return res


def run(cmd, **kwargs):
    assert not cmd.startswith("sudo")
    return _run(cmd, **kwargs)


def run_as_root(cmd, **kwargs):
    # For commands prefixed with sudo: interestingly there's no need to set
    # interactive=True, as sudo writes the prompt straight on the terminal
    # device, see sudo(8), option '-S'.
    if os.geteuid() == 0:
        rootcmd = cmd
    else:
        rootcmd = "sudo " + cmd
    return _run(rootcmd, **kwargs)


def write_file_as_root(target, content, mode="0644"):
    """
    Write content to a file, as root. If the file already exists,
    it is overwritten.

    'content' might be a string or a list. If it's a list, it is
    assumed that each element of the list corresponds to a line,
    and doesn't end with a newline.

    If content is not empty, and it doesn't end with a newline,
    then a newline is added.

    Implementation-wise, the content is written to a temporary file
    at first, and then it's moved to the target location with sudo.

    NOTES for unit testing: make sure to test empty list, non-empty
    list, empty string, non-empty string.
    """
    if isinstance(content, list):
        if len(content) > 0:
            content = "\n".join(content) + "\n"
        else:
            content = ""
    elif len(content) > 0 and not content.endswith("\n"):
        content += "\n"

    with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
        f.write(content)

    cmd = f"install -m {mode} {f.name} {target}"
    res = run_as_root(cmd, interactive=True, log=False)
    os.remove(f.name)
    if res.returncode != 0:
        raise RuntimeError(f"Couldn't write file '{target}'")


def is_pkg_installed(pkg):
    res = run(f"dpkg -s {pkg} 2>/dev/null | grep -q 'ok installed'")
    return res.returncode == 0


def apt_update(force=False):
    """
    Run apt-get update.

    Not as straightforward as expected, as apt-get update returns
    zero even in case of failure.  So we have to check the output
    in order to figure out if something went wrong. If there's an
    error (ie. if any line starts with 'Err:') we fail. If there's
    a warning (line starting with 'W:') we keep going.

    There are some warnings that we want to ignore, for example
    « Key is stored in legacy trusted.gpg keyring » is a well-known
    occurrence with third-party repos, and we really don't care.

    This function might be called many times, but we don't need to
    run it all the time, so optimize a bit and run once per minute
    at most, unless force is True.
    """
    if "last" not in apt_update.__dict__:
        apt_update.last = 0

    if not force:
        now = int(time.time())
        if now - apt_update.last < 60:
            return

    say("Updating APT cache")

    cmd = "apt-get update"
    if os.geteuid() != 0:
        cmd = "sudo " + cmd

    print("> $ " + cmd)

    process = subprocess.Popen(
        cmd, shell=True, stderr=subprocess.STDOUT, stdout=subprocess.PIPE, text=True
    )

    apt_failed = False
    for line in process.stdout:
        sys.stdout.write(line)
        sys.stdout.flush()
        if line.startswith("Err:"):
            apt_failed = True
    process.wait()

    if process.returncode != 0 or apt_failed is True:
        raise RuntimeError("Couldn't update the APT cache")

    apt_update.last = int(time.time())


def apt_upgrade():
    """
    Run apt-get full-upgrade, ask for confirmation.
    """
    say("Upgrading the system")

    cmd = "apt-get full-upgrade"
    res = run_as_root(cmd, interactive=True)
    if res.returncode != 0:
        raise RuntimeError("Couldn't upgrade the system")


def apt_install_pkgs(pkgs, reinstall=False, yes=True):
    """
    Install packages.
    """
    apt_update()
    say_install_packages(pkgs)
    cmd = "apt-get install"
    if reinstall:
        cmd += " --reinstall"
    if yes:
        cmd += " --yes"
    if isinstance(pkgs, list):
        pkgs = " ".join(pkgs)
    cmd += f" {pkgs}"
    res = run_as_root(cmd, interactive=True)
    if res.returncode != 0:
        raise RuntimeError(f"Couldn't install package(s) '{pkgs}'")


def apt_remove_pkgs(pkgs):
    """
    Remove packages.
    """
    say_remove_packages(pkgs)
    cmd = "apt-get -y --purge remove"
    if isinstance(pkgs, list):
        pkgs = " ".join(pkgs)
    cmd += f" {pkgs}"
    res = run_as_root(cmd, interactive=True)
    if res.returncode != 0:
        raise RuntimeError(f"Couldn't remove package(s) '{pkgs}'")


def apt_list_upgradable_pkgs():
    """
    List packages that can be upgraded.

    Note that, ideally, we shouldn't parse apt output, apt says not
    to do that.  So we could use apt-cache instead, something like:
    "apt-cache -qq show '?upgradable'". But the output is super
    verbose.
    """
    apt_update()
    cmd = "apt list -qq --upgradable"
    res = run(cmd)
    if res.returncode != 0:
        raise RuntimeError("Couldn't list upgradable package(s)")
    pkgs = []
    for line in res.stdout.splitlines():
        parts = line.split("/", maxsplit=1)
        if len(parts) == 1:
            continue
        pkgs.append(parts[0])
    return pkgs


def _get_metapackages(search):
    """Return an array of (pkg, description, installed)"""
    out = []

    # Run the search command
    res = run(search)
    if res.returncode != 0:
        return out

    # For every result
    for line in res.stdout.splitlines():
        try:
            # Example: kali-linux-default - Kali Linux default system
            pkg, desc = line.split(" - ", maxsplit=1)
        except ValueError:
            continue
        installed = is_pkg_installed(pkg)
        out += [(pkg, desc, installed)]

    return out


def apt_list_kali_metapackages():
    """
    List the Kali metapackages of interest.

    We return a list of tuples: (pkg, desc, installed) where:
    * 'pkg' is the package name
    * 'desc' is the package description
    * 'installed' is True if the package is installed, False otherwise
    """
    apt_update()
    tools = _get_metapackages(
        "apt-cache search --names-only ^kali-tools | grep -v ' menu$'"
    )
    meta = _get_metapackages(
        "apt-cache search --names-only ^kali-linux"
        r" | grep -vE '\-(arm|core|nethunter) '"
    )
    return tools + meta


def install_program(program):
    say_install_program(program)
    dst = program
    src = get_datafile(os.path.basename(program))
    cmd = f"install -v -D -m0755 {src} {dst}"
    res = run_as_root(cmd, interactive=True)
    if res.returncode != 0:
        raise RuntimeError(f"Couldn't install program {src} -> {dst}")


def is_program_installed(program):
    if not os.path.isfile(program):
        return False
    if not os.access(program, os.X_OK):
        return False
    src = get_datafile(os.path.basename(program))
    return filecmp.cmp(src, program)


def enable_service_now(service):
    say_enable_service(service)
    cmd = f"systemctl start {service}"
    res = run_as_root(cmd, interactive=True)
    if res.returncode != 0:
        raise RuntimeError(f"Couldn't start service '{service}'")
    cmd = f"systemctl enable {service}"
    res = run_as_root(cmd, interactive=True)
    if res.returncode != 0:
        raise RuntimeError(f"Couldn't enable service '{service}'")


def disable_service_now(service):
    say_disable_service(service)
    cmd = f"systemctl stop {service}"
    res = run_as_root(cmd, interactive=True)
    if res.returncode != 0:
        raise RuntimeError(f"Couldn't stop service '{service}'")
    cmd = f"systemctl disable {service}"
    res = run_as_root(cmd, interactive=True)
    if res.returncode != 0:
        raise RuntimeError(f"Couldn't disable service '{service}'")


def is_service_up(service):
    cmd = f"systemctl is-active {service}"
    res = run(cmd)
    if res.stdout.strip() != "active":
        return False
    cmd = f"systemctl is-enabled {service}"
    res = run(cmd)
    if res.stdout.strip() != "enabled":
        return False
    return True


# -------- SHELL HELPERS -------- #


def shell_from_shell_path(shell_path):
    """Get a shell name from a shell path, ie: /bin/zsh -> zsh"""
    for sh in ["zsh", "bash", "fish", "ksh", "sh", "tcsh"]:
        if f"bin/{sh}" in shell_path:
            return sh
    return None


def shell_path_from_shell(shell):
    """Get a shell path from a shell name, ie: zsh -> /usr/bin/zsh"""
    if shell == "bash":
        return "/usr/bin/bash"
    elif shell == "zsh":
        return "/usr/bin/zsh"
    else:
        return None


def get_configured_login_shell():
    """Get the login shell, as configured in /etc/passwd"""
    username = getpass.getuser()
    res = run(f"getent passwd {username}")
    if res.returncode != 0:
        logger.debug("getent passwd %s: %s", username, res.returncode)
        return None

    line = res.stdout.strip()
    _, _, shell_path = line.rpartition(":")
    # per passwd(5), the last field might be empty,
    # in such case it defaults to bin/sh
    if not shell_path:
        shell_path = "/bin/sh"
    logger.debug("Detected shell_path: %s (from /etc/passwd)", shell_path)

    shell = shell_from_shell_path(shell_path)
    logger.debug("Corresponding shell: %s", shell)
    return shell


def get_effective_login_shell():
    """Get the current login shell, as given by the 'SHELL' variable"""
    shell_path = os.getenv("SHELL")
    if shell_path is None:
        logger.debug("shell_path is None")
        return None
    logger.debug("Detected shell_path: %s (from SHELL)", shell_path)

    shell = shell_from_shell_path(shell_path)
    logger.debug("Corresponding shell: %s", shell)
    return shell


def install_shell_if_missing(shell):
    """Install shell if missing"""
    if is_pkg_installed(shell):
        return
    pkgs = [shell]
    if shell == "zsh":
        pkgs += ["zsh-autosuggestions", "zsh-syntax-highlighting"]
    apt_install_pkgs(pkgs)


def change_login_shell(sh):
    """Change the login shell (requires logout+login to be effective)"""
    shell_path = shell_path_from_shell(sh)
    if shell_path is None:
        raise NotImplementedError(f"Can't set user shell to: {sh}")
    install_shell_if_missing(sh)
    say_change_login_shell(sh)
    cmd = f"chsh -s {shell_path}"
    res = run(cmd, interactive=True)
    if res.returncode != 0:
        msg = f"Couldn't set login shell to '{sh}' ({shell_path})\n"
        msg += "Check console output for details."
        raise RuntimeError(msg)


def get_shell_config_files():
    """Return a list of tuples (default, current)"""
    configs = []
    homedir = str(Path.home())
    skeldir = "/etc/skel"
    for rc in [".bashrc", ".zshrc"]:
        default = os.path.join(skeldir, rc)
        current = os.path.join(homedir, rc)
        if not os.path.isfile(default):
            continue
        if not os.access(default, os.R_OK):
            continue
        configs += [(default, current)]
    return configs


def apply_shell_config_files(configs):
    """Take a list of tuples from get_shell_config_files()"""
    say("Resetting shell config files")
    try:
        for default, current in configs:
            if os.path.isfile(current):
                backup = backup_filename(current)
                print(f"> '{current}' -> '{backup}'")
                shutil.copyfile(current, backup)
            print(f"> '{default}' -> '{current}'")
            shutil.copyfile(default, current)
    except OSError as e:
        raise RuntimeError(str(e))


def _get_shell_rc_file(shell):
    rc_files = {
        "bash": ".bashrc",
        "zsh": ".zshrc",
        "pwsh": ".config/powershell/Microsoft.PowerShell_profile.ps1",
    }
    try:
        rc_file = rc_files[shell]
    except KeyError:
        logger.debug("unsupported shell %s", shell)
        return None
    path = str(Path.home() / rc_file)
    if not os.path.isfile(path):
        logger.debug("%s is not a file", path)
        return None
    if not os.access(path, os.R_OK):
        logger.debug("no read access to %s", path)
        return None
    return path


def _get_kali_config_block(data):
    """The whole config block as lines, including delimiters"""
    start_delimiter = "# START KALI CONFIG VARIABLES"
    stop_delimiter = "# STOP KALI CONFIG VARIABLES"
    res = []
    keep = False
    for line in data.splitlines():
        if line == start_delimiter:
            res = [line]
            keep = True
        elif line == stop_delimiter:
            res += [line]
            keep = False
            break
        elif keep:
            res += [line]

    if keep is True:
        # there was a start delimiter, but no stop delimiter
        return []

    return res


def _get_shell_config_from_file(filename):
    """
    Retrieve the config values that are within the Kali block.

    This input is supposed to be a UNIX shell config file.
    Powershell is not supported.
    """
    # Get the configuration block
    with open(filename) as f:
        data = f.read()
    content = _get_kali_config_block(data)
    if not content:
        logger.debug("config is empty")
        return None

    # Get all variables assignments and store it in a dictionary.
    # For simplicity, we expect values to consist only in alphanumeric chars,
    # underscore or dash. If it's not the case, the variable is discarded.
    config = {}
    for line in content:
        # drop empty lines and comments
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        # parse line
        key, sep, value = line.partition("=")
        # there's should be one equal sign
        if not sep:
            continue
        # key should be a valid shell variable name
        if not re.fullmatch(r"[a-zA-Z_][a-zA-Z0-9_]*", key):
            continue
        # let's be restrictive here: we only expect keywords anyway,
        # so we accept only alphanumeric plus underscore and dash.
        if value.startswith('"') and value.endswith('"'):
            value = value[1:-1]
        elif value.startswith("'") and value.endswith("'"):
            value = value[1:-1]
        if not re.fullmatch(r"[a-zA-Z0-9_\-]*", value):
            continue
        # let's keep this line!
        config[key] = value

    return config


def _apply_shell_config_to_file(filename, config):
    """Apply some config values to the shell config file"""
    # Get the configuration block
    with open(filename) as f:
        data = f.read()
    content = _get_kali_config_block(data)
    if not content:
        return

    # Keep a copy for later
    content_orig = content.copy()

    # Set some config (add it if missing)
    for key, value in config.items():
        print(f"> Setting: {key}='{value}'")

        changed = False
        for idx, line in enumerate(content):
            line = line.strip()
            if not line.startswith(key):
                continue
            content[idx] = f"{key}='{value}'"
            changed = True
            break

        if changed:
            continue

        content.append(f"{key}='{value}'")

    # Build a text block from content
    config_block_orig = "\n".join(content_orig)
    config_block = "\n".join(content)

    logger.debug("config orig:\n%s", config_block_orig)
    logger.debug("config     :\n%s", config_block)

    # Return if no changes
    if config_block_orig == config_block:
        return

    # Do the actual replacement
    data = data.replace(config_block_orig, config_block)

    # Write to file
    tmp = tmp_filename(filename)
    with open(tmp, "w") as f:
        f.write(data)

    backup = backup_filename(filename)
    shutil.copyfile(filename, backup)
    shutil.move(tmp, filename)


def _apply_config_to_shells(config, shells):
    """Apply a configuration to a list of shells.

    The configuration is a simple (key, value) dict to be applied "as is".
    """

    if not config:
        return

    for shell in shells:
        rc_file = _get_shell_rc_file(shell)
        if not rc_file:
            continue
        print(f"> Updating: {rc_file}")
        _apply_shell_config_to_file(rc_file, config)


def get_current_shell_config():
    """
    Retrieve the config for the current shell. Beware that config can
    be an *empty dict* (config file found, but no config variables in
    there), which is different from *None* (no config file found).
    """
    shell = get_effective_login_shell()
    if not shell:
        logger.debug("No shell detected")
        return None
    logger.debug("Detected shell: %s", shell)

    rc_file = _get_shell_rc_file(shell)
    if not rc_file:
        logger.debug("No shell_rc detected")
        return None
    logger.debug("Detected shell_rc: %s", rc_file)

    # Get shell config
    unix_conf = _get_shell_config_from_file(rc_file)
    if not unix_conf:
        return None

    # Translate config
    config = {}
    if "PROMPT_ALTERNATIVE" in unix_conf:
        val = unix_conf["PROMPT_ALTERNATIVE"]
        config["prompt-style"] = val
    if "NEWLINE_BEFORE_PROMPT" in unix_conf:
        val = True if unix_conf["NEWLINE_BEFORE_PROMPT"] == "yes" else False
        config["prompt-settings"] = {}
        config["prompt-settings"]["newline"] = val

    return config


def apply_shell_config(config):
    """Apply the config to the supported shells"""
    say("Configuring command prompt")

    unix_conf = {}
    powershell_conf = {}

    # Translate config
    if "prompt-style" in config:
        key = "PROMPT_ALTERNATIVE"
        val = config["prompt-style"]
        unix_conf[key] = val
        powershell_conf[f"${key}"] = val
    if "prompt-settings" in config:
        settings = config["prompt-settings"]
        if "newline" in settings:
            key = "NEWLINE_BEFORE_PROMPT"
            val = "yes" if settings["newline"] is True else "no"
            unix_conf[key] = val
            powershell_conf[f"${key}"] = val

    # Apply config
    _apply_config_to_shells(unix_conf, ["bash", "zsh"])
    _apply_config_to_shells(powershell_conf, ["pwsh"])
