#!/usr/bin/python3

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

import logging
import os
import shutil
import sys
import termios

from snack import (
    Grid,
    SnackScreen,
    Textbox,
)

from .settings.aptrepositories import AptRepositoriesSetting
from .settings.openssl import OpenSSLSetting
from .settings.samba import SambaSetting
from .settings.ssh import SSHSetting
from .settings.sysctl import SysctlSetting
from .settings.virtualization import VirtualizationSetting
from .toolkit import (
    KaliCheckboxBar,
    KaliConfirmWindow,
    KaliInfoWindow,
    KaliListboxWindow,
    KaliRadioBar,
    KaliSimpleCheckboxTree,
    KaliTextboxFullSize,
    KaliWindow,
)
from .utils import (
    apply_shell_config,
    apply_shell_config_files,
    apt_install_pkgs,
    apt_list_kali_metapackages,
    apt_list_upgradable_pkgs,
    apt_remove_pkgs,
    apt_upgrade,
    bold,
    change_login_shell,
    get_configured_login_shell,
    get_current_shell_config,
    get_effective_login_shell,
    get_shell_config_files,
    is_pkg_installed,
)

LOG_FILE = "kali-tweaks-log.txt"
LOG_LEVEL = logging.DEBUG

if "DEBUG" in os.environ:
    # Can't use stdout while the UI is on, so we send logs to a file instead
    logging.basicConfig(filename=LOG_FILE, filemode="w", level=LOG_LEVEL)

logger = logging.getLogger("kali-tweaks")

HEIGHT = 12


# -------- UI HELPERS -------- #

SCREEN = None


class KaliTweaksScreen:
    def __init__(self):
        self.screen = None
        self.terminal_size = (-1, -1)

    def create(self):
        """
        Create a SnackScreen instance.
        """
        if self.screen is not None:
            logger.warning("The screen object already exists, can't create a new one")
            return

        self.screen = SnackScreen()
        # Theming for checkbox and radiobox
        self.screen.setColor("CHECKBOX", "black", "lightgray")
        self.terminal_size = shutil.get_terminal_size()
        logger.debug("Screen created: %d x %d", self.screen.width, self.screen.height)

    def destroy(self):
        """
        Destroy the current SnackScreen instance.
        """
        if self.screen is None:
            logger.warning("The screen object doesn't exist, can't destroy it")
            return

        self.screen.finish()
        self.screen = None
        self.terminal_size = (-1, -1)
        logger.debug("Screen destroyed")

    def prepare(self):
        """
        Prepare the current screen before showing a window. In other words,
        recreate it if the dimensions of the terminal have changed.
        """
        if not self.screen:
            self.create()
            return

        if shutil.get_terminal_size() != self.terminal_size:
            self.destroy()
            self.create()

    def show_info_window(self, level, msg):
        """
        Show an informative message in a window.

        level can be: info, warning, error.
        """
        self.prepare()
        KaliInfoWindow(self.screen, HEIGHT, level, msg)

    def show_confirm_window(self, title, msg, ok_label=None, cancel_label=None):
        """
        Show a confirmation window, for user to say yes or no.

        Returns the button that was chosen by user.
        """
        keywords = {"title": title}
        if ok_label:
            keywords["ok_label"] = ok_label
        if cancel_label:
            keywords["cancel_label"] = cancel_label
        self.prepare()
        button = KaliConfirmWindow(self.screen, HEIGHT, msg, **keywords)
        return button

    def show_window(self, title, content, buttons):
        """
        Show an interactive window, for user to change settings.

        Returns the button that was chosen by user.
        """
        self.prepare()
        window = KaliWindow(self.screen, title, content, buttons)
        button = window.run_once()
        return button

    def show_listbox_window(self, title, items, buttons):
        """
        Show a list in a window. I guess this should be merged with
        show_window() at some point.
        """
        self.prepare()
        button, line = KaliListboxWindow(self.screen, HEIGHT, title, items, buttons)
        return button, line


def wait_enter_keystroke(message="> Press Enter to continue..."):
    termios.tcflush(sys.stdin, termios.TCIFLUSH)
    input(bold(message))


def _highlight(title, message, title_color="blue", wait_enter=True):
    """
    Display a message prominently on the terminal.

    Unless wait_enter=False, wait for use to press <Enter> before going
    any further.  This is to make sure that users have time to read the
    message.
    """
    colors = {
        "red": "\033[1;31m",
        "green": "\033[1;32m",
        "yellow": "\033[1;33m",
        "blue": "\033[1;34m",
    }
    reset = "\033[00m"
    color = colors.get(title_color, reset)
    print()
    print(f"┏━({color}{title}{reset})")
    for line in message.splitlines():
        print(f"┃ {line}")
    print("┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    print()
    if wait_enter:
        wait_enter_keystroke()


def highlight_message(message, wait_enter=True):
    title = "Message from Kali developers"
    _highlight(title, message, wait_enter=wait_enter)


def highlight_error(message, wait_enter=True):
    title = "Something went wrong!"
    _highlight(title, message, wait_enter=wait_enter, title_color="red")


def say_reload_terminal(first_line):
    msg = f"{first_line}\n"
    msg += "Please close and restart all your terminal instances"
    msg += " for the change to take effect."
    highlight_message(msg)


def say_log_out(first_line):
    msg = f"{first_line}\n"
    msg += "Please log out and log back in"
    msg += " for the change to take effect."
    highlight_message(msg)


# -------- CLASSES -------- #


class KaliGenericContent(Grid):
    def __init__(self, settings):
        # assume vertical padding, hence 'rows = n_settings + 1'
        super().__init__(1, len(settings) + 1)

        self._orig_config = settings
        self._widgets = []

        # No setting to show?
        if not settings:
            msg = r"No setting in this section ¯\_(ツ)_/¯"
            widget = KaliTextboxFullSize(HEIGHT, msg)
            self.setField(widget, 0, 0, growx=1)
            return

        # Iterate on settings and make widgets
        used_height = 0
        for item in settings:
            heading = item["desc"]
            if heading[-1] not in [".", ":"]:
                heading += ":"
            choices = item["choices"]
            value = item["value"]
            widget = None
            if item["type"] == "single-choice":
                widget = KaliRadioBar(heading, choices, value)
            elif item["type"] == "multiple-choices":
                widget = KaliCheckboxBar(heading, choices, value)
            else:
                raise NotImplementedError

            self._widgets.append(widget)
            used_height += widget.title_height + len(choices)

        # Add the widgets
        for index, widget in enumerate(self._widgets):
            self.setField(widget, 0, index, growx=1)

        # Add vertical padding
        padding = HEIGHT - used_height
        if padding <= 0:
            return
        widget = Textbox(1, padding, "")
        self.setField(widget, 0, len(self._widgets), growx=1)

    def getOrigSelection(self):
        """
        Return the original selection
        """
        data = {}
        for setting in self._orig_config:
            setting_name = setting["name"]
            data[setting_name] = setting.value
        return data

    def getCurrentSelection(self):
        """
        Return the current selection, eg:

        {
          "my-radiobar-setting": "selected-item",
          "my-checkboxbar-setting": {
            "setting1": True,
            "setting2": False,
          }
        }
        """
        data = {}
        for setting, widget in zip(self._orig_config, self._widgets):
            setting_name = setting["name"]
            data[setting_name] = widget.getSelection()
        return data

    def getChanges(self):
        """
        Return only what changed between the current selection and
        the original selection.
        """
        changes = {}
        for setting, widget in zip(self._orig_config, self._widgets):
            setting_name = setting["name"]
            orig = setting["value"]
            curr = widget.getSelection()
            if curr is None:
                # 'None' here means that the value was initially unknown,
                # and was not modified by user, hence it's still None.
                continue
            elif type(curr) is str:
                if orig is None:
                    orig = ""
                assert type(orig) is str
                if curr != orig:
                    changes[setting_name] = curr
            elif type(curr) is dict:
                if orig is None:
                    orig = {}
                assert type(orig) is dict
                ch = {}
                for key in curr:
                    if key not in orig:
                        ch[key] = curr[key]
                    elif curr[key] != orig[key]:
                        ch[key] = curr[key]
                if ch:
                    changes[setting_name] = ch
            else:
                raise NotImplementedError

        logger.debug("changes: %s", changes)

        return changes


# -------- SECONDARY SCREENS -------- #


def do_shell_change_default_login():
    title = "Default Login Shell"
    buttons = [
        ("Apply", "ok"),
        ("Back", "cancel"),
    ]

    # Get effective and configured login shells, make sure there's no change pending
    configured_shell = get_configured_login_shell()
    effective_shell = get_effective_login_shell()
    if configured_shell != effective_shell:
        msg = f"Your shell is configured to be {configured_shell},"
        msg += f" however at the moment it is effectively {effective_shell}.\n"
        msg += "It seems that you have changes pending,"
        msg += " please log out and log back in for these changes to take effect."
        SCREEN.show_info_window("warning", msg)
        return

    # Prepare a config
    config = {}
    config["shell"] = effective_shell

    # All settings
    all_settings = []

    setting = {
        "name": "shell",
        "desc": "Default Login Shell to use",
        "type": "single-choice",
        "choices": [
            (
                "Bash",
                "Bourne Again SHell",
                "bash",
            ),
            (
                "ZSH",
                "Z SHell",
                "zsh",
            ),
        ],
        "value": config["shell"],
    }
    all_settings.append(setting)

    # Show the window with all the settings
    content = KaliGenericContent(all_settings)
    btn = SCREEN.show_window(title, content, buttons)
    if btn != "ok":
        return

    # Get the settings that changed
    changes = content.getChanges()
    if changes == {}:
        return

    # Ask for confirmation
    old_shell = config["shell"]
    new_shell = changes["shell"]
    msg = ""
    msg += f"Your current login shell is {old_shell}.\n"
    msg += f"Changing it to {new_shell}.\n"
    msg += "Proceed?"

    btn = SCREEN.show_confirm_window(title, msg)
    if btn != "ok":
        return

    # Do the actual job
    SCREEN.destroy()
    try:
        change_login_shell(new_shell)
    except RuntimeError as e:
        highlight_error(str(e))
        return

    # One last notification
    say_log_out("Your login shell was changed!")


def do_shell_reset():
    configs = get_shell_config_files()

    title = "Reset Shell Config"
    msg = "The following config files will be replaced by Kali's default:\n"
    for default, current in configs:
        msg += f"* {current} (replaced with {default})\n"
    msg += "WARNING: Your changes in those files, if any, will be lost.\n"
    msg += "Proceed?"

    btn = SCREEN.show_confirm_window(title, msg)
    if btn != "ok":
        return

    # Do the actual job
    SCREEN.destroy()
    try:
        apply_shell_config_files(configs)
        wait_enter_keystroke()
    except RuntimeError as e:
        highlight_error(str(e))
        return

    # One last notification
    say_reload_terminal("Your shell config files were reset!")


def do_shell_change_prompt():
    title = "Command Prompt"
    buttons = [
        ("Apply", "ok"),
        ("Back", "cancel"),
    ]

    # Get the current shell config
    config = get_current_shell_config()

    # Couldn't get the config?
    if config is None:
        msg = "Failed to parse shell configuration!\n"
        msg += "This is expected if your shell configuration files are"
        msg += " too old, and all you need to do is to reset it to the"
        msg += " latest default.\n"
        msg += "Do you want to reset your shell configuration files?"

        btn = SCREEN.show_confirm_window(title, msg)
        if btn == "ok":
            return "go-to-reset"
        else:
            return

    # All settings
    all_settings = []

    setting = {
        "name": "prompt-style",
        "desc": "Prompt Style",
        "type": "single-choice",
        "choices": [
            (
                "Two Lines",
                "Dual blue line prompt",
                "twoline",
            ),
            (
                "One Line",
                "Single blue line prompt",
                "oneline",
            ),
            (
                "BackTrack",
                "Legacy BackTrack red prompt",
                "backtrack",
            ),
        ],
        "value": config.get("prompt-style", "twoline"),
    }
    all_settings.append(setting)

    setting = {
        "name": "prompt-settings",
        "desc": "Prompt Settings",
        "type": "multiple-choices",
        "choices": [
            (
                "Newline",
                "Add a new line between output and prompt",
                "newline",
            ),
        ],
        "value": config.get("prompt-settings", True),
    }
    all_settings.append(setting)

    # Show the window with all the settings
    content = KaliGenericContent(all_settings)
    btn = SCREEN.show_window(title, content, buttons)
    if btn != "ok":
        return

    # Get the settings that changed
    changes = content.getChanges()
    if changes == {}:
        return

    # Ask for confirmation
    msg = ""
    if "prompt-style" in changes:
        new = changes["prompt-style"]
        msg += f"Setting prompt to '{new}'.\n"
    if "prompt-settings" in changes:
        new = changes["prompt-settings"]
        if "newline" in new:
            if new["newline"] is True:
                msg += "Enabling new line before prompt.\n"
            else:
                msg += "Disabling new line before prompt.\n"
    msg += "Proceed?"

    btn = SCREEN.show_confirm_window(title, msg)
    if btn != "ok":
        return "retry"

    # Do the actual job
    SCREEN.destroy()
    try:
        apply_shell_config(changes)
    except RuntimeError as e:
        highlight_error(str(e))
        return

    # One last notification
    say_reload_terminal("Your shell config files were updated!")


# -------- MAIN SCREENS -------- #


def do_shell_screen():
    title = "Shell & Prompt"
    items = [
        ("Configure Prompt", "Configure the command prompt"),
        ("Default Login Shell", "Set the default login shell"),
        ("Reset Shell Config", "Reset the shell config files to their default"),
    ]
    buttons = [
        ("Select", "ok"),
        ("Back", "cancel"),
    ]
    ok_funcs = [
        do_shell_change_prompt,
        do_shell_change_default_login,
        do_shell_reset,
    ]

    while True:
        btn, line = SCREEN.show_listbox_window(title, items, buttons)
        if btn != "ok":
            break
        func = ok_funcs[line]
        while True:
            ret = func()
            if ret == "go-to-reset":
                func = do_shell_reset
                ret = func()
            if ret != "retry":
                break


def do_virtual_screen():
    title = "Virtualization"
    buttons = [
        ("Select", "ok"),
        ("Back", "cancel"),
    ]

    # Detect what kind of virtualization is that
    virt = VirtualizationSetting()
    virt.detect()
    logger.debug("Virtualization: %s", virt.type)

    # Not virtualized?
    if virt.type is None:
        msg = "Couldn't detect virtualization type."
        SCREEN.show_info_window("info", msg)
        return

    # Virtualized, but we don't provide any additional config?
    if not virt.is_configurable():
        msg = (
            f"Virtualization detected: {virt.pretty_name}.\n\n"
            "Nothing to configure for this virtualization technology."
        )
        SCREEN.show_info_window("info", msg)
        return

    # We might want to configure, or to unconfigure, or to do nothing,
    # in case the system is configured and there's no way to unconfigure.
    if virt.needs_configuration():
        items = [("Configure", virt.config_desc)]
        action = virt.configure
    elif virt.is_unconfigurable():
        items = [("Unconfigure", virt.undo_desc)]
        action = virt.unconfigure
    else:
        msg = (
            f"Virtualization detected: {virt.pretty_name}.\n\n"
            "Your system is already configured for this virtualization technology."
        )
        SCREEN.show_info_window("info", msg)
        return

    # Display the one-liner menu
    btn, _ = SCREEN.show_listbox_window(title, items, buttons)
    if btn != "ok":
        return

    # Do the actual job
    SCREEN.destroy()
    try:
        action()
    except RuntimeError as e:
        highlight_error(str(e))
        return

    # Good-bye message
    msg = (
        f"For more information about Kali inside {virt.pretty_name}, please refer to:\n"
        f"{virt.help_url}"
    )
    highlight_message(msg)


def do_network_repo_screen():
    title = "Network Repositories"
    buttons = [
        ("Apply", "ok"),
        ("Back", "cancel"),
    ]

    # All settings
    all_settings = []
    config = AptRepositoriesSetting().load()

    setting = {
        "name": "extra-repos",
        "desc": "Additional Kali repositories",
        "type": "multiple-choices",
        "choices": [
            (
                "bleeding-edge",
                "Automatically packaged and potentially unstable",
                "kali-bleeding-edge",
            ),
            (
                "experimental",
                "Staging area for work-in-progress packages",
                "kali-experimental",
            ),
        ],
        "value": config["extra-repos"],
    }
    all_settings.append(setting)

    setting = {
        "name": "mirror",
        "desc": "Mirrors",
        "type": "single-choice",
        "choices": [
            (
                "Cloudflare",
                "Download packages from the Cloudflare CDN",
                "kali.download",
            ),
            (
                "Community",
                "Download packages from a community mirror",
                "http.kali.org",
            ),
        ],
        "value": config["mirror"],
    }
    all_settings.append(setting)

    setting = {
        "name": "protocol",
        "desc": "Protocol",
        "type": "single-choice",
        "choices": [
            (
                "HTTP",
                "Use the HTTP protocol",
                "http",
            ),
            (
                "HTTPS",
                "Use the HTTPS protocol",
                "https",
            ),
        ],
        "value": config["protocol"],
    }
    all_settings.append(setting)

    # Show the window with all the settings
    content = KaliGenericContent(all_settings)
    btn = SCREEN.show_window(title, content, buttons)
    if btn != "ok":
        return

    # Get the settings that changed
    changes = content.getChanges()
    if changes == {}:
        return

    # Ask for confirmation
    msg = ""
    if "extra-repos" in changes:
        new = changes["extra-repos"]
        for repo, enabled in new.items():
            if enabled:
                msg += f"Enabling {repo} repository.\n"
            else:
                msg += f"Disabling {repo} repository.\n"
    if "mirror" in changes:
        new = changes["mirror"]
        msg += f"Setting mirror to '{new}'.\n"
    if "protocol" in changes:
        new = changes["protocol"]
        msg += f"Setting protocol to '{new}'.\n"
    msg += "Proceed?"

    btn = SCREEN.show_confirm_window(title, msg)
    if btn != "ok":
        return "retry"

    # Do the actual job
    SCREEN.destroy()
    try:
        AptRepositoriesSetting().save(changes)
    except RuntimeError as e:
        highlight_error(str(e))
        return

    # Good-bye message
    msg = ""
    if "extra-repos" in changes:
        msg += (
            "For more information about Kali repositories, please refer to:\n"
            "https://www.kali.org/docs/general-use/kali-apt-sources/\n"
        )
    if (
        "extra-repos" in changes
        and changes["extra-repos"].get("bleeding-edge", False) is True
    ):
        msg += (
            "For more information about the Kali Bleeding Edge branch:\n"
            "https://www.kali.org/docs/general-use/kali-bleeding-edge/\n"
        )
    if "mirror" in changes:
        # I don't think we have doc regarding community vs cloudflare mirror?
        pass
    if "protocol" in changes and changes["protocol"] == "https":
        msg += (
            "For more information about HTTPS support:\n"
            "https://www.kali.org/blog/kali-linux-repository-https-support/\n"
        )
    if msg:
        highlight_message(msg)
    else:
        wait_enter_keystroke()


def do_hardening_screen():
    title = "Hardening Settings"
    buttons = [
        ("Apply", "ok"),
        ("Back", "cancel"),
    ]

    # All hardening settings
    all_settings = []

    # Kernel setting
    kernel_setting = {
        "name": "kernel",
        "desc": "Kernel settings:",
        "type": "multiple-choices",
        "choices": [
            (
                "Restrict dmesg",
                "Forbid unprivileged users to run dmesg",
                "dmesg-restrict",
            ),
            (
                "Privileged Ports",
                "Restrict ports < 1024 for privileged users",
                "ports-restrict",
            ),
        ],
        "value": SysctlSetting().load(),
    }
    all_settings.append(kernel_setting)

    # Wide compatibility setting
    compat_setting = {
        "name": "wide-compat",
        "desc": (
            "In Wide Compatibility mode, old protocols, ciphers and algorithms are enabled,\n"
            "allowing access to legacy services. Uncheck to select Strong Security instead."
        ),
        "type": "multiple-choices",
        "choices": [],
        "value": {},
    }
    all_settings.append(compat_setting)

    if is_pkg_installed("openssl"):
        try:
            config = OpenSSLSetting().load()
        except (KeyError, IOError) as e:
            msg = (
                "Failed to read the OpenSSL configuration!\n\n"
                "Please make sure that your Kali Linux installation"
                " is up to date and try again.\n\n"
                f"{e}"
            )
            SCREEN.show_info_window("error", msg)
            return
        choice = (
            "OpenSSL",
            " ",
            "openssl",
        )
        compat_setting["choices"].append(choice)
        compat_setting["value"]["openssl"] = False
        if config["hardening"] == "compat":
            compat_setting["value"]["openssl"] = True

    if is_pkg_installed("smbclient"):
        try:
            config = SambaSetting().load()
        except RuntimeError as e:
            msg = "Failed to read the Samba configuration!\n\n" f"{e}"
            SCREEN.show_info_window("error", msg)
            return
        choice = (
            "Samba client",
            " ",
            "samba",
        )
        compat_setting["choices"].append(choice)
        compat_setting["value"]["samba"] = False
        if config["hardening"] == "compat":
            compat_setting["value"]["samba"] = True

    if is_pkg_installed("openssh-client"):
        try:
            config = SSHSetting().load()
        except RuntimeError as e:
            msg = "Failed to read the SSH configuration!\n\n" f"{e}"
            SCREEN.show_info_window("error", msg)
            return
        choice = (
            "SSH client",
            " ",
            "ssh",
        )
        compat_setting["choices"].append(choice)
        compat_setting["value"]["ssh"] = False
        if config["hardening"] == "compat":
            compat_setting["value"]["ssh"] = True

    # Show the window with all the settings
    content = KaliGenericContent(all_settings)
    btn = SCREEN.show_window(title, content, buttons)
    if btn != "ok":
        return

    # Get the settings that changed
    changes = content.getChanges()
    if changes == {}:
        return

    # Ask for confirmation
    msg = ""
    if "kernel" in changes:
        new = changes["kernel"]
        if "dmesg-restrict" in new:
            if new["dmesg-restrict"] is True:
                msg += "Restricting dmesg to privileged users.\n"
            else:
                msg += "Allowing dmesg for unprivileged users.\n"
        if "ports-restrict" in new:
            if new["ports-restrict"] is True:
                msg += "Restricting ports < 1024 to privileged users.\n"
            else:
                msg += "Allowing unprivileged to use any ports.\n"
    if "wide-compat" in changes:
        new = changes["wide-compat"]
        if "openssl" in new:
            new_val = new["openssl"]
            msg += f"Setting wide compatibility for OpenSSL to '{new_val}'.\n"
        if "samba" in new:
            new_val = new["samba"]
            msg += f"Setting wide compatibility for Samba to '{new_val}'.\n"
        if "ssh" in new:
            new_val = new["ssh"]
            msg += f"Setting wide compatibility for SSH to '{new_val}'.\n"
    msg += "Proceed?"

    btn = SCREEN.show_confirm_window(title, msg)
    if btn != "ok":
        return "retry"

    # Do the actual job
    SCREEN.destroy()
    msg = ""
    err_msg = ""

    if "kernel" in changes:
        new = changes["kernel"]
        try:
            SysctlSetting().save(new)
            msg += (
                "For more information about Kernel configuration, please refer to:\n"
                "https://www.kali.org/docs/general-use/kernel-configuration/\n"
            )
        except RuntimeError as e:
            err_msg += "Failed to update kernel settings!\n\n" f"{e}\n"

    # XXX the quick way to avoid re-indenting it all...
    changes = changes.get("wide-compat", {})
    if "openssl" in changes:
        if changes["openssl"] is True:
            new_config = {"hardening": "compat"}
        else:
            new_config = {"hardening": "secure"}
        try:
            OpenSSLSetting().apply(new_config)
            msg += (
                "For more information about OpenSSL configuration, please refer to:\n"
                "https://www.kali.org/docs/general-use/openssl-configuration/\n"
            )
        except (KeyError, IOError, ValueError) as e:
            err_msg += (
                "Failed to update the OpenSSL configuration!\n\n"
                "Please make sure that your Kali Linux installation"
                " is up to date and try again.\n\n"
                f"{e}\n"
            )
    if "samba" in changes:
        if changes["samba"] is True:
            new_config = {"hardening": "compat"}
        else:
            new_config = {"hardening": "secure"}
        try:
            SambaSetting().apply(new_config)
            msg += (
                "For more information about Samba configuration, please refer to:\n"
                "https://www.kali.org/docs/general-use/samba-configuration/\n"
            )
        except RuntimeError as e:
            err_msg += "Failed to update Samba configuration!\n\n" f"{e}\n"
    if "ssh" in changes:
        if changes["ssh"] is True:
            new_config = {"hardening": "compat"}
        else:
            new_config = {"hardening": "secure"}
        try:
            SSHSetting().apply(new_config)
            msg += (
                "For more information about SSH configuration, please refer to:\n"
                "https://www.kali.org/docs/general-use/ssh-configuration/\n"
            )
        except RuntimeError as e:
            err_msg += "Failed to update SSH configuration!\n\n" f"{e}\n"

    # Good-bye message
    needs_wait = True
    if msg:
        highlight_message(msg)
        needs_wait = False
    if err_msg:
        highlight_error(err_msg)
        needs_wait = False
    if needs_wait:
        wait_enter_keystroke()


def do_metapackages_screen():
    title = "Metapackages"
    buttons = [
        ("Apply", "ok"),
        ("Back", "cancel"),
    ]

    # Get metapackages - destroy the screen, as we know that before
    # listing metapackage we might run 'apt update', which might ask
    # for sudo password, hence requires access to the terminal.
    SCREEN.destroy()
    try:
        items = apt_list_kali_metapackages()
    except RuntimeError as e:
        highlight_error(str(e))
        return
    if not items:
        msg = "No Kali Tools found. Are you really running Kali Linux? Or do you have an apt issue?"
        SCREEN.show_info_window("info", msg)
        return

    # Tweak the package names for better display
    new_items = []
    for pkg, desc, installed in items:
        if pkg.startswith("kali-tools-"):
            pkg = pkg.removeprefix("kali-tools-")
            desc = desc.removeprefix("Kali Linux ")
        new_items += [(pkg, desc, installed)]
    items = new_items

    # Show the window with all the tools
    content = KaliSimpleCheckboxTree(HEIGHT, items)
    btn = SCREEN.show_window(title, content, buttons)
    if btn != "ok":
        return

    # Sum up changes
    selection = content.getSelection()
    to_install = []
    to_remove = []
    for pkg, _, installed in items:
        package = pkg
        if not package.startswith("kali-linux-"):
            package = "kali-tools-" + package
        if pkg in selection and not installed:
            to_install += [package]
        elif pkg not in selection and installed:
            to_remove += [package]

    # No change?
    if to_install == [] and to_remove == []:
        return

    # Ask for confirmation
    msg = ""
    if to_install:
        msg += "Packages to install: "
        msg += ", ".join(to_install)
        msg += "\n"
    if to_remove:
        msg += "Packages to remove: "
        msg += ", ".join(to_remove)
        msg += "\n"
    msg += "Proceed?"

    btn = SCREEN.show_confirm_window(title, msg)
    if btn != "ok":
        return "retry"

    # Do the actual job

    # Step 1. Propose to upgrade the system first.
    #
    # It only makes sense if user wants to install packages. This is optional,
    # but recommended, so we gently push our users in this direction.
    #
    # If they agree, we won't ask for confirmation in step 2, as step 1 will
    # take a while, and we don't want to break and wait after that, that would
    # be bad UX.
    #
    # The whole point of asking for user confirmation is that APT might do
    # something wrong (ie. remove a ton of packages), so we want to give a
    # chance to review the changes before acting. But if step 1 is agreed on,
    # then there are very little risk that APT does something wrong while
    # installing the metapackages, so it seems acceptable not to ask.
    #
    # Cf. #38.

    SCREEN.destroy()
    say_yes = False
    to_upgrade = []

    if to_install:
        try:
            to_upgrade = apt_list_upgradable_pkgs()
        except RuntimeError as e:
            highlight_error(str(e))
            return

    if to_upgrade:
        msg = (
            "We detected that your system is not fully up-to-date. "
            "Before installing a new metapackage, we will perform a "
            "full upgrade. If you don't want this, please select <Skip>."
        )
        btn = SCREEN.show_confirm_window(title, msg, cancel_label="Skip")
        SCREEN.destroy()
        if btn == "ok":
            try:
                apt_upgrade()
            except RuntimeError as e:
                highlight_error(str(e))
                return
            say_yes = True

    # Step 2. Install or remove metapckages
    try:
        if to_remove:
            apt_remove_pkgs(to_remove)
        if to_install:
            apt_install_pkgs(to_install, yes=say_yes)
    except RuntimeError as e:
        highlight_error(str(e))
        return

    # Good-bye message
    msg = (
        "For more information about Kali's metapackages, please refer to:\n"
        "https://www.kali.org/docs/general-use/metapackages/"
    )
    highlight_message(msg)


def do_main_screen():
    title = "Main Menu"
    items = [
        ("Hardening", "Configure the system for extra security"),
        ("Metapackages", "Install specific subsets of tools for particular needs"),
        ("Network Repositories", "Configure network repositories for APT sources"),
        ("Shell & Prompt", "Configure the shell and command prompt"),
        ("Virtualization", "Additional configurations for Virtual Machines"),
    ]
    buttons = [
        ("Select", "ok"),
        ("Quit", "cancel"),
    ]
    ok_funcs = [
        do_hardening_screen,
        do_metapackages_screen,
        do_network_repo_screen,
        do_shell_screen,
        do_virtual_screen,
    ]

    while True:
        btn, line = SCREEN.show_listbox_window(title, items, buttons)
        if btn != "ok":
            break
        func = ok_funcs[line]
        while True:
            ret = func()
            if ret != "retry":
                break


def main():
    global SCREEN
    SCREEN = KaliTweaksScreen()
    try:
        do_main_screen()
    finally:
        SCREEN.destroy()
        # Print logs to stdout, if any
        if os.path.isfile(LOG_FILE):
            with open(LOG_FILE, "r") as f:
                print(f.read(), end="")
            os.remove(LOG_FILE)


if __name__ == "__main__":
    main()
