#!/usr/bin/env python3

import os
import re
import sys
from argparse import ArgumentError
from itertools import chain
from pathlib import Path
from socketserver import BaseRequestHandler, TCPServer
from subprocess import DEVNULL, PIPE, Popen, check_output, run

PORT = "7777"
REMOTE_HOST = "localhost"
REMOTE_PORT = "7777"
AUTH_SOCK_DIR_PATTERN = re.compile(r"ssh-")
AUTH_SOCK_PATTERN = re.compile(f"agent.[0-9]+$")


class ReusableTCPServer(TCPServer):
    allow_reuse_address = True


def check_ssh_auth_sock(**kwargs):
    completed = run(
        ["ssh-add", "-l"],
        stdout=PIPE,
        stderr=DEVNULL,
        universal_newlines=True,
        **kwargs,
    )
    return (
        completed.returncode == 0
        or completed.stdout.rstrip() == "The agent has no identities."
    )


def get_ssh_auth_sock():
    if "SSH_AUTH_SOCK" in os.environ and check_ssh_auth_sock():
        return os.environ["SSH_AUTH_SOCK"]
    for q in chain.from_iterable(
        p.iterdir()
        for p in Path(os.environ.get("TMPDIR") or "/tmp").iterdir()
        if p.is_dir() and AUTH_SOCK_DIR_PATTERN.match(p.name)
    ):
        q_str = str(q)
        if (
            q.is_socket()
            and AUTH_SOCK_PATTERN.fullmatch(q.name)
            and check_ssh_auth_sock(env={**os.environ, "SSH_AUTH_SOCK": q_str})
        ):
            return q_str
    print("starting new ssh-agent")
    return check_output(
        ["bash", "-c", '. <(ssh-agent); echo "${SSH_AUTH_SOCK}"'],
        universal_newlines=True,
    )


def create_tcp_server(server_address, data):
    if not isinstance(data, bytes):
        data = data.encode()

    class RequestHandler(BaseRequestHandler):
        def handle(self):
            self.request.sendall(data)

    return ReusableTCPServer(server_address, RequestHandler)


def usage():
    print(f"{sys.argv[0]} [-h|--help] [--sudo] ssh-args script")
    print("\t-h, --help  display this help message")
    print("\t--sudo      run remote script as root")
    print("\tssh-args    arguments passed to ssh; include a hostname here")
    print("\tscript      script to run remotely")


def main():
    args = sys.argv[1:-1]
    if len(args) < 1:
        raise ArgumentError(None, "No host specified")
    if "--help" in args or "-h" in args:
        usage()
        exit(0)
    if "--sudo" in args:
        sudo = "sudo "
        args.remove("--sudo")
    else:
        sudo = ""
    environ = {**os.environ, "SSH_AUTH_SOCK": get_ssh_auth_sock()}
    run(["ssh-add", "id_ed25519"], stderr=DEVNULL, env=environ)
    with open(sys.argv[-1], "rb") as fh:
        with create_tcp_server(("127.0.0.1", int(PORT)), fh.read()) as server:
            sp = Popen(
                [
                    "ssh",
                    "-R",
                    ":".join((PORT, REMOTE_HOST, REMOTE_PORT)),
                    "-o",
                    "ForwardX11=no",
                    "-At",
                    *args,
                    f"{sudo}bash -c '. <(cat </dev/tcp/{REMOTE_HOST}/{REMOTE_PORT})'",
                ],
                env=environ,
            )
            server.handle_request()
            sp.communicate()


if __name__ == "__main__":
    main()