#!/usr/bin/env python3

# This script can be used in order to automatically refresh the
# installation of ZTP Server to the latest code on GitHub.  This can
# be useful in order to pull bug fixes or run the latest version of
# various development branches.

# INSTRUCTIONS:
#  - The script must be run by a user with sudo access.

import argparse
import os
import re
import sys

BACKUP_FOLDER = "/tmp/ztps_backup"
BACKUP_CONFIG_FOLDER = "/tmp/ztps_backup/home"
GLOBAL_CONFIG_FILE = "/etc/ztpserver/ztpserver.conf"

GIT_REPO_FOLDER = "/tmp/ztps"


def create_folder(path):
    os.system(f"sudo rm -fr '{path}'")
    os.system(f"sudo mkdir {path}")
    os.system(f"sudo chmod 777 {path}")


def all_files(path):
    result = []
    for top, dirs, files in os.walk(path):
        result += [os.path.join(top, d) for d in dirs]
        result += [os.path.join(top, f) for f in files]

    return result


def refresh_ztps(home_folder, branch):
    # Backup configuration files
    create_folder(BACKUP_FOLDER)
    create_folder(BACKUP_CONFIG_FOLDER)

    os.system(f"sudo cp {GLOBAL_CONFIG_FILE} {'%s/ztpserver.conf' % BACKUP_FOLDER}")
    os.system(f"sudo cp -r {home_folder}/* {BACKUP_CONFIG_FOLDER}")

    # Uninstall
    paths = all_files("/usr/lib")
    paths += all_files("/usr/local")
    for path in [p for p in paths if re.match(".*ztpserver.*", p)]:
        os.system(f"sudo rm -rf {path}")

    # Clone repository
    os.system(f"sudo rm -rf '{GIT_REPO_FOLDER}'")
    result = os.system(
        f"""
        sudo git clone https://github.com/arista-eosplus/ztpserver {GIT_REPO_FOLDER};
        cd {GIT_REPO_FOLDER};
        sudo git checkout {branch};
        sudo git pull;

        sudo python setup.py build;
        sudo python setup.py install
        """
    )
    if result != 0:
        sys.exit(f"ERROR: failed to re-install ZTP Server (branch {branch})")

    # Reconfigure
    os.system(f"sudo cp {'%s/ztpserver.conf' % BACKUP_FOLDER} {GLOBAL_CONFIG_FILE}")

    os.system(
        f"sudo cp {BACKUP_CONFIG_FOLDER}/bootstrap/bootstrap.conf {home_folder}/bootstrap/bootstrap.conf"
    )
    os.system(f"sudo cp {BACKUP_CONFIG_FOLDER}/neighbordb' {home_folder}/neighbordb")

    files_1 = [y for y in [x.replace(home_folder + "/", "") for x in all_files(home_folder)] if y]
    files_2 = [
        y
        for y in [
            x.replace(BACKUP_CONFIG_FOLDER + "/", "") for x in all_files(BACKUP_CONFIG_FOLDER)
        ]
        if y
    ]

    for file_path in [x for x in files_2 if x not in files_1]:
        try:
            os.system(f"sudo cp {BACKUP_CONFIG_FOLDER}/{file_path} {home_folder}/{file_path}")
        except IOError:
            # directory
            create_folder(f"{home_folder}/{file_path:f}")

    # Set permissions
    os.system(f"sudo chmod 777 {home_folder} -R")
    os.system(f"sudo chmod 777 {GLOBAL_CONFIG_FILE} -R")

    # Clean up
    os.system(f"sudo rm -fr '{BACKUP_FOLDER}'")
    os.system(f"sudo rm -fr '{GIT_REPO_FOLDER}'")
    print("\nDone!")


def main():
    parser = argparse.ArgumentParser(prog="refresh_ztps")
    parser.add_argument(
        "-b",
        "--branch",
        metavar="BRANCH",
        nargs="?",
        default="master",
        help="ZTP Server branch to be used for the refresh",
        action="store",
    )
    parser.add_argument(
        "-f",
        "--home-folder",
        metavar="HOME_FOLDER",
        nargs="?",
        default="/usr/share/ztpserver",
        help="ZTP Server home folder",
        action="store",
    )
    args = parser.parse_args()
    refresh_ztps(args.home_folder, args.branch)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass
