install

#!/usr/bin/env python3

import os
import pathlib
import platform
import sys

from subprocess import run, PIPE

COLOR_BLUE = "\033[94m"
COLOR_RED = "\033[31m"
COLOR_GREEN = "\033[32m"
COLOR_YELLOW = "\033[93m"
COLOR_NONE = "\033[00m"
ICON_CHECK_MARK = "\U00002705"

# Use environ instead of getenv since we want it to fail if var not found
HOME_DIR = os.environ["HOME"]
CURRENT_DIR = os.getcwd()
ZSH_DIR = os.path.join(HOME_DIR, ".oh-my-zsh")
ZSH_THEME_DIR = os.path.join(ZSH_DIR, "custom/themes")
ZSH_PLUGIN_DIR = os.path.join(ZSH_DIR, "custom/plugins")

    
INSTALL_FILE_DICT = {
    "config/bash/.bashrc": ".bashrc",
    "config/zsh/.zshenv": ".zshenv",
    "config/zsh/.zshrc": ".zshrc",
    "config/.vimrc": ".vimrc",
}
# Convert INSTALL_FILE_DICT to be absolute paths
# Keys are assumed are relative to current directory, values are assumed to be in HOME
INSTALL_FILE_DICT = {os.path.join(CURRENT_DIR, s): os.path.join(HOME_DIR, d) for s, d in INSTALL_FILE_DICT.items()}

INSTALL_DIR_DICT = {
    "bin": "bin"
}
# Convert INSTALL_DIR_DICT to be absolute paths
# Keys are assumed are relative to current directory, values are assumed to be in HOME
INSTALL_DIR_DICT = {os.path.join(CURRENT_DIR, s): os.path.join(HOME_DIR, d) for s, d in INSTALL_DIR_DICT.items()}


def __print_colored(message: str, color: str):
    """
    Function that simply prints the message but in a specific color provided by the param.
    :param message: the message to be printed
    :param color: the color to use for the message
    """
    print(f"{color}{message}{COLOR_NONE}")


def _print_info(message):
    """
    Function that prints a message with INFO colors (NONE).
    :param message: the message to be printed
    """
    __print_colored(message, COLOR_NONE)


def _print_title(message):
    """
    Function that prints a message with TITLE colors (BLUE).
    :param message: the message to be printed
    """
    __print_colored(message, COLOR_BLUE)


def _print_done(message):
    """
    Function that prints a message with DONE colors (GREEN).
    :param message: the message to be printed
    """
    __print_colored(message, COLOR_GREEN)


def _print_warning(message):
    """
    Function that prints a message with WARNING colors (YELLOW).
    :param message: the message to be printed
    """
    __print_colored(message, COLOR_YELLOW)


def _print_error(message):
    """
    Function that prints a message with ERROR colors (RED).
    :param message: the message to be printed
    """
    __print_colored(message, COLOR_RED)


def __create_symlink(source: str, destination: str):
    """
    Function that creates a symlink from source to destination.
    :param source: source file or directory that needs to be linked.
    :param destination: destination file or directory where the link will be made.
    """
    if __is_destination_already_set(source, destination):
        _print_info(f"Symlink {destination} -> {source} already exists, skipping.")
        return

    try:
        os.symlink(source, destination)
        _print_done(f"Symlink {destination} -> {source} created")
    except OSError as err:
        _print_error(f'Failed to create symlink from src: {source} dest: {destination}. Reason: {err}')


def __prompt_user_yes_no(question: str) -> bool:
    """
    Function that prompts user the question from the argument for a Yes/No answer.
    :param question: str, question the user will be prompted for.
    :return: bool, True if user answers Yes, False otherwise.
    """
    while (answer := input(f"{question} (y/n): ").lower()) not in {"y", "yes", "n", "no"}:
        _print_warning("Invalid input.")

    if answer[0] == "y":
        return True

    return False


def _execute(command: list):
    """
    Function that executes subprocess.run command and captures the output.
    :param command: command in form of a list
    :return bool, True if command succeded, false otherwise.
    """
    result = run(command, stdout=PIPE, stderr=PIPE)
    
    if result.returncode != 0:
        return False
    
    return True


def __is_destination_already_set(source: str, destination: str):
    """
    Function that checks if the destination is already symlinked to the source.
    :param source: source file or directory that needs to be linked.
    :param destination: destination file or directory where the link will be made.
    """
    if os.path.islink(destination) and os.readlink(destination) == source:
        return True

    return False


def check_if_dest_exists():
    """
    Function that checks if any of the files from INSTALL_FILE_DICT exists,
    or if any of the directories from INSTALL_DIR_DICT exists, and moves them using `.backup` suffix.
    """
    for source, dest in INSTALL_DIR_DICT.items():
        if not __is_destination_already_set(source, dest):
            if os.path.isdir(dest):
                _print_info(f"Directory {dest} found, backing up")
                os.rename(dest, f"{dest}.backup")

    for source, dest in INSTALL_FILE_DICT.items():
        if not __is_destination_already_set(source, dest):
            if os.path.isfile(dest):
                _print_info(f"File {dest} found, backing up")
                os.rename(dest, f"{dest}.backup")


def do_setup_zsh():
    """
    Function to clone oh-my-zsh to $HOME/.oh-my-zsh
    Clone powerlevel10k theme to ZSH_CUSTOM/custom/themes
    Clone zsh-autosuggestions, zsh-completions, zsh-history-substring-search, zsh-syntax-highlighting in ZSH_CUSTOM/custom/plugins
    """
    if not os.path.isdir(ZSH_DIR):
        res = _execute(f"git clone https://github.com/ohmyzsh/ohmyzsh.git --depth 5 --single-branch {ZSH_DIR}".split())
        if res:
            _print_done(f"ZSH dir: {ZSH_DIR} clone complete.")
        else:
            _print_error(f"ZSH dir: {ZSH_DIR} clone failed.")
    else:
        _print_info(f"ZSH dir: {ZSH_DIR} already exists, skipping.")
    
    for plugin in ["zsh-autosuggestions", "zsh-completions", "zsh-history-substring-search", "zsh-syntax-highlighting"]:
        if not os.path.isdir(f"{ZSH_PLUGIN_DIR}/{plugin}"):
            res = _execute(f"git clone https://github.com/zsh-users/{plugin}.git --depth 5 --single-branch {ZSH_PLUGIN_DIR}/{plugin}".split())
            if res:
                _print_done(f"ZSH plugin: {plugin} clone complete.")
            else:
                _print_error(f"ZSH plugin: {plugin} clone failed.")
        else:
            _print_info(f"ZSH plugin: {plugin} already exists, skipping.")
    
    if not os.path.isdir(f"{ZSH_THEME_DIR}/powerlevel10k"):
        res = _execute(f"git clone https://github.com/romkatv/powerlevel10k.git --depth 5 --single-branch {ZSH_THEME_DIR}/powerlevel10k".split())
        if res:
            _print_done(f"ZSH theme: powerlevel10k clone complete.")
        else:
            _print_error(f"ZSH theme: powerlevel10k clone failed.")
    else:
        _print_info(f"ZSH theme: powerlevel10k already exists, skipping.")

    if os.path.isfile("stop-the-warnings.diff"):
        cwd = os.getcwd()
        _print_info("Patching oh-my-zsh ...")
        res = _execute(f"cd {ZSH_DIR} && git apply {cwd}/stop-the-warnings.diff && cd {cwd}".split())
        if res:
            _print_done(f"Patching ZSH complete.")
        else:
            _print_error(f"Patching ZSH failed.")


def do_symlinks():
    """
    Function that performs all the needed symlinks for files and directories.
    """
    for s, d in INSTALL_DIR_DICT.items():
        __create_symlink(s, d)

    for s, d in INSTALL_FILE_DICT.items():
        __create_symlink(s, d)

    p10k_file = os.path.join(HOME_DIR, ".p10k.zsh")
    install_p10k = True
    if os.path.isfile(p10k_file):
        response = __prompt_user_yes_no(f"\nConfig {p10k_file} found, do you want to replace it?")
        if response:
            os.rename(p10k_file, f"{p10k_file}.backup")
        else:
            install_p10k = False
            _print_info("Skipping .p10k.zsh, file already exists and wasn't replaced.")

    if install_p10k:
        while (answer := input(f"""\nPowerlevel10k options,
            \n    1) lean
            \n    2) pretty
            \n(1/2): """).lower()) not in {"1", "2"}: _print_warning("Invalid input.")

        if answer[0] == "1":
            __create_symlink(os.path.join(CURRENT_DIR, "config/zsh/.p10k-lean.zsh"), os.path.join(HOME_DIR, ".p10k.zsh"))
        else:
            __create_symlink(os.path.join(CURRENT_DIR, "config/zsh/.p10k-pretty.zsh"), os.path.join(HOME_DIR, ".p10k.zsh"))


def configure_git():
    """
    Function that configures git.
    Some things like name are hardcoded, but sensitive things like email are prompted.
    GPG key is looked up based on email entered.
    """
    gitconfig_file = os.path.join(HOME_DIR, ".gitconfig")
    configure_git_config = True
    if os.path.isfile(gitconfig_file):
        response = __prompt_user_yes_no(f"Config {gitconfig_file} found, do you want to overwrite it?")
        if response:
            os.rename(gitconfig_file, f"{gitconfig_file}.backup")
        else:
            configure_git_config = False
            _print_info("Skipping .gitconfig, file already exists and wasn't replaced.")
    
    if configure_git_config:
        git_alias_file = os.path.join(CURRENT_DIR, "config/gitconfig.aliases")
        _execute(["git", "config", "--global", "user.name", "Ankit Sadana"])
        
        _execute("git config --global http.sslVerify true".split())
        _execute("git config --global core.editor vim".split())
        _execute("git config --global init.defaultBranch main".split())
        _execute(f"git config --global include.path {git_alias_file}".split())
        
        _execute("git config --global gpg.program gpg".split())
        _execute("git config --global commit.gpgsign true".split())

        email = input("Enter your git associated email: ")
        _execute(f"git config --global user.email {email}".split())

        _print_info(f"Looking for gpg key for {email} ...")
        result = run(f"gpg --list-secret-keys --keyid-format LONG {email}".split(),
            stdout=PIPE, stderr=PIPE)

        if result.returncode != 0:
            _print_warning("No GPG key found")
            return

        key = result.stdout.decode("utf-8").splitlines()[0].split()[1].split("/")[1]
        _execute(f"git config --global user.signingkey {key}".split())
        _print_done("GPG key found, installed to gitconfig")


def configure_gpg():
    """
    Function that configures gpg agent to be used with pinentry
    """
    gpg_directory = pathlib.Path("~/.gnupg").expanduser()
    if not gpg_directory.is_dir():
        gpg_directory.mkdir(mode=0o700)

    conf_file = gpg_directory / "gpg.conf"
    conf_file.write_text("use-agent\n", encoding="utf-8", newline="\n")

    agent_file = gpg_directory / "gpg-agent.conf"
    if platform.system() == "Darwin":
        agent_file.write_text("pinentry-program /opt/homebrew/bin/pinentry-mac\n", encoding="utf-8", newline="\n")
    else:
        agent_file.write_text("pinentry-program /usr/bin/pinentry-tty\n", encoding="utf-8", newline="\n")


def main():
    """
    Function that does all the orchestration between functions.
    TODO: Configure this into a menu, with auto option to go through expected flow.
    """
    check_if_dest_exists()

    _print_title("\nSetting up oh-my-zsh, plugins and themes")
    do_setup_zsh()

    _print_title("\nCreating symlinks")
    do_symlinks()

    _print_title("\nConfiguring gpg")
    configure_gpg()

    _print_title("\nConfiguring git-config")
    configure_git()

    _print_info(f"\nAll done! {ICON_CHECK_MARK}")


if __name__ == "__main__":
    main()