#!/usr/bin/env -S python3 -u
"""
Adapted from:
https://github.com/theupdateframework/python-tuf/tree/f8deca31ccea22c30060f259cb7ef2588b9c6baa/examples/client
"""


import argparse
import inspect
import json
import os
import re
import socketserver
import sys
from urllib import request

from tuf.ngclient import Updater


def parse_args():
    parser = argparse.ArgumentParser()
    for parameter in inspect.signature(main).parameters.values():
        if parameter.name.startswith("_"):
            continue
        if "KEYWORD" in parameter.kind.name:
            parser.add_argument(
                "--" + parameter.name.replace("_", "-"),
                default=parameter.default,
            )
    return parser.parse_args()


def semver(s):
    return tuple(s.lstrip("v").split("."))


def name_matches(name, f):
    return re.match(name, f)


def readline():
    result = []
    c = sys.stdin.read(1)
    while c != "\n":
        result.append(c)
        c = sys.stdin.read(1)
    result.append(c)
    return "".join(result)


class Handler(socketserver.BaseRequestHandler):
    def __init__(self, *args, tuf_server=None, updater=None, **kwargs):
        self.tuf_server = tuf_server
        self.updater = updater
        super().__init__(*args, **kwargs)

    def handle(self):
        self.request.settimeout(10)
        os.dup2(self.request.fileno(), sys.stdin.fileno())
        os.dup2(self.request.fileno(), sys.stdout.fileno())

        print("Welcome to the firmware update admin console!")
        print("What type of firmware would you like to download from the TUF server?")
        print(
            "Whichever type you pick, we will pull the latest version from the server."
        )
        print("Types:")
        with request.urlopen(f"{self.tuf_server}/targets.json") as response:
            targets = json.load(response)
        all_target_files = list(targets["signed"]["targets"].keys())
        print("-", "\n- ".join({file.split("_")[0] for file in all_target_files}))

        print("Enter type name: ")
        name = readline().strip()
        if "." in name:
            # People were trying to bypass our version check with regex tricks! Not allowed!
            print("Not allowed!")
            return
        filenames = list(
            sorted(
                [f for f in all_target_files if name_matches(name, f)],
                key=lambda s: semver(s),
            )
        )
        if len(filenames) == 0:
            print("Sorry, file not found!")
            return
        filename = filenames[-1]

        print(f"Downloading {filename}")

        info = self.updater.get_targetinfo(filename)
        if info is None:
            print("Sorry, file not found!")
            return

        with open("/dev/urandom", "rb") as f:
            name = f.read(8).hex()
        path = self.updater.download_target(
            info,
            filepath=f"/tmp/{name}.{os.path.basename(info.path)}",
        )
        os.chmod(path, 0o755)

        print(f"Running {filename}")
        child = os.fork()
        if child == 0:
            os.execl(path, path)
        else:
            os.wait()
            os.remove(path)


def main(tuf_server="http://localhost:8001", server_port="8002", **_):
    repo_metadata_dir = "/tmp/tuf_server_metadata"
    if not os.path.isdir(repo_metadata_dir):
        if os.path.exists(repo_metadata_dir):
            raise RuntimeError(
                f"{repo_metadata_dir} already exists and is not a directory"
            )
        os.mkdir(repo_metadata_dir)
        with request.urlopen(f"{tuf_server}/root.json") as response:
            root = json.load(response)
        with open(f"{repo_metadata_dir}/root.json", "w") as f:
            json.dump(root, f, indent=2)

    updater = Updater(
        metadata_dir=repo_metadata_dir,
        metadata_base_url=tuf_server + "/metadata/",
        target_base_url=tuf_server + "/targets/",
    )
    updater.refresh()

    def return_handler(*args, **kwargs):
        return Handler(*args, **kwargs, tuf_server=tuf_server, updater=updater)

    print("Running server")
    with socketserver.ForkingTCPServer(
        ("0", int(server_port)), return_handler
    ) as server:
        server.serve_forever()


if __name__ == "__main__":
    main(**parse_args().__dict__)