#!/usr/bin/env python3

import argparse as ap
import base64 as b64
import os
import datetime as dt
import socket as s
import subprocess as subp
import sys
import tempfile
import time
import threading

import requests as r

def utcnow():
    return dt.datetime.now(dt.timezone.utc).timestamp()

def ssh_sign(message, key_path, namespace="udp-holepunch"):
    """Sign a message using ssh-keygen."""
    return subp.check_output(["ssh-keygen", "-f", key_path, "-Y", "sign", "-n", namespace], input=message, text=True, stderr=subp.DEVNULL)

def ssh_verify(message, signature, allowed_keys, namespace="udp-holepunch"):
    """Verify a message using ssh-keygen.

    `allowed_keys` should be a list of keys of the form

      ssh-whatever encodedpublickeyhere
    """
    allowed_keys_fix = []
    principals = set()
    for k in allowed_keys:
        k = k.strip()
        if k.count(' ') < 2:
            ## No principal
            i = 0
            while "p%d" % i in principals:
                i += 1
            principals.add("p%d" % i)
            allowed_keys_fix.append(("p%d " % i) + k)
        else:
            a, b, principal = k.split(' ')
            principals.add(principal)
            allowed_keys_fix.append(' '.join((principal, a, b)))

    with tempfile.TemporaryDirectory() as td:
        valid_f = os.path.join(td, "valid")
        with open(valid_f, 'w') as f:
            f.write('\n'.join(allowed_keys_fix))

        signature_f = os.path.join(td, "signature")
        with open(signature_f, 'w') as f:
            f.write(signature)

        for p in principals:
            ## Inexplicably, check_call doesn't have `input`
            try:
                subp.check_output(["ssh-keygen", "-f", valid_f, "-Y", "verify", "-n", namespace, "-s", signature_f, "-I", p], input=message, text=True, stderr=subp.DEVNULL)
            except subp.CalledProcessError:
                continue
            else:
                break
        else:
            return False

    return True

def udp_holepunch_loop(sock, other_addr, no_ack=False, timeout=60):
    """Finalize the UDP holepunching.

    The general idea here is to blast them with "hello!" packets; once they receive one from us,
    they reply with a "done!" packet, and we do the same. With `no_ack`, we just blast a couple
    "hello!" packets and then return.
    """
    if no_ack:
        ## Send a couple anyways but don't listen for a response
        for _ in range(5):
            sock.sendto(b"hello!", other_addr)
    else: 
        found = False

        start = time.monotonic()
        sock.settimeout(0.5)
        while True:
            sock.sendto(b"hello!" if not found else b"done!", other_addr)

            try:
                resp = sock.recvfrom(100)
            except TimeoutError:
                if time.monotonic() - start > timeout:
                    return False

                continue

            if resp[1] != other_addr:
                continue
            
            resp = resp[0]
            if resp == b"done!":
                sock.sendto(b"done!", other_addr)
                break

            if not found:
                found = True
                continue

        sock.settimeout(0.5)
        while True:
            try:
                resp = sock.recvfrom(100)
            except TimeoutError:
                break

    return True

def udp_holepunch_server(local_port, server_port, server_path, no_ack=False, allowed_keys=None, timeout=60, sig_valid=60*10):
    """Serve as the UDP holepunching "server".

    Poll the given `server_path`; if a message is found from `udp_holepunch_client`, begin the
    holepunching process.
    """
    other_ip = None
    other_port = None

    ## Poll the server_path for a file telling us what to phone
    while True:
        resp = r.get(server_path, stream=True)
        if resp.status_code == 404:
            return None

        server_ip = resp.raw._connection.sock.getpeername()[0]

        info = resp.text.strip()
        if allowed_keys is not None:
            if ';' not in info:
                ## Missing signature
                raise Exception("Missing signature on the webserver")

            info, sig = info.split(';', 1)
            sig = b64.b64decode(sig).decode("utf-8")
            if not ssh_verify(info, sig, allowed_keys):
                ## Invalid signature
                raise Exception("Invalid signature on the webserver")

        if ',' not in info:
            raise Exception("Missing time on the webserver")

        info, other_time = info.split(',')
        other_time = float(other_time)

        if abs(utcnow() - other_time) > sig_valid:
            raise Exception("Signature expired on webserver")

        other_ip, other_port = info.split(':')
        break

    other_port = int(other_port)

    ## Send to the intermediate server
    c = s.socket(s.AF_INET, s.SOCK_DGRAM)
    c.settimeout(0.5)
    c.bind(("", local_port))

    if server_port == 0:
        server_port = other_port

    ## Reasonable number of tries
    for _ in range(3):
        c.sendto(b"hello!", (server_ip, server_port))
        time.sleep(0.1)

    internal_port = c.getsockname()[1]

    ## Send to the client directly
    other_addr = (other_ip, other_port)

    if not udp_holepunch_loop(c, other_addr, no_ack=no_ack, timeout=timeout):
        raise Exception("Holepunching timeout")

    c.close()

    ## We have established holepunching
    return (internal_port, other_ip, other_port)

def parse_packet(data):
    """Parse packet from tcpdump."""
    data = data.decode("utf-8").strip()
    source_info = data.split(' ')[2].split('.')
    source_ip = '.'.join(source_info[:-1])
    source_port = source_info[-1]

    return source_ip, int(source_port)

def udp_holepunch_client(local_port, server_name, server_port, server_path, no_ack=False, key_path=None, timeout=60):
    """Serve as the UDP holepunching "client".

    Post a message to `server_path` on `server_name` via SSH, then listen using `tcp_dump` to find
    our own external port and the "server"'s external port.
    """
    server_ip = None

    ## Setup the socket
    c = s.socket(s.AF_INET, s.SOCK_DGRAM)
    c.bind(("", local_port))
    c.settimeout(0.5)

    ## Internal port number downstream applications must use
    internal_port = c.getsockname()[1]

    auto_server_port = server_port == 0
    if auto_server_port:
        server_port = internal_port

    ## Our external info
    external_ip = None
    external_port = None

    ## Info for the other end
    other_ip = None
    other_port = None

    lock = threading.RLock()

    def server_ops():
        ## Functions run on the intermediate server

        nonlocal external_ip
        nonlocal external_port
        nonlocal other_ip
        nonlocal other_port
        nonlocal server_ip

        with lock:
            ## We do the following on the server:
            ## 1. Get our own external IP via $SSH_CONNECTION;
            ## 2. Clear out the previous poll file
            ## 3. Find our external port
            ## 4. Write our information to the poll file
            ## 5. Listen for the server's UDP ping
            ## 6. Remove the poll file
            proc = subp.Popen(
                [
                    "ssh", "-tt", server_name,
                    """
echo $SSH_CONNECTION
rm {server_path}
sudo tcpdump -n -c 1 -i eth0 udp port {server_port}
head -n 1 > {server_path}
SERVER_PORT=`head -n 1`
sudo tcpdump -c 1 -n -i eth0 udp port $SERVER_PORT
rm {server_path}
                    """.strip().replace('\n', ';').format(server_path=server_path, server_port=server_port)
                ],
                stdout=subp.PIPE,
                stderr=subp.DEVNULL,
                stdin=subp.PIPE,
            )

            ## 1. Get the IP of the webserver
            ssh_conn = proc.stdout.readline().strip().split()
            server_ip = ssh_conn[-2].decode("utf-8")

        ## 2. Clear out the previous poll file (nop)

        ## 3. Find our external port
        ## Clear tcpdump trash at start, and SSH might complain about TTY
        while True:
            i = proc.stdout.readline()
            if i.startswith(b"tcpdump"):
                break
        proc.stdout.readline()

        data = proc.stdout.readline()
        with lock:
            external_ip, external_port = parse_packet(data)

        proc.stdout.readline()
        proc.stdout.readline()
        proc.stdout.readline()

        ## 4. Write our information to the poll file
        info_string = "%s:%s,%d" % (external_ip, external_port, utcnow())
        if key_path is not None:
            ## Add a signature
            sig = ssh_sign(info_string, key_path)
            sig = b64.b64encode(sig.encode("utf-8")).decode("utf-8")
            info_string += ';' + sig
        info_string += '\n'

        proc.stdin.write(info_string.encode("utf-8"))
        proc.stdin.flush()
        ## Head prints to stdout
        proc.stdout.readline()

        ## 5. Listen for the server's UDP ping

        ## Write the server port to the intermediate
        proc.stdin.write(
            ("%d\n" % (external_port if auto_server_port else server_port)).encode("utf-8")
        )
        proc.stdin.flush()
        ## Head prints to stdout
        proc.stdout.readline()

        ## Wait for things to settle
        time.sleep(1)

        ## Clear tcpdump trash at start
        data = proc.stdout.readline()
        data = proc.stdout.readline()

        while True:
            data = proc.stdout.readline()
            if not data:
                continue
            ip, port = parse_packet(data)

            if ip == external_ip:
                ## Ignore our own packets
                continue

            with lock:
                other_ip = ip
                other_port = int(port)
                proc.communicate()
                break

        ## 6. Remove the poll file (nop)

    ## Run operations on the intermediate server
    server_thread = threading.Thread(target=server_ops, daemon=True)
    server_thread.start()

    ## 3. Find our external port
    ## Send UDP packets to the intermediate server
    while True:
        with lock:
            if external_port is not None:
                break
            if server_ip is None:
                ## FIXME: Should probably be a semaphore before the while loop
                time.sleep(0.1)
                continue
        c.sendto(b"hello!", (str(server_ip), server_port))
        time.sleep(0.5)

    server_thread.join()

    ## Now we're done with the intermediate: talk directly to the real server
    other_addr = (other_ip, other_port)

    udp_holepunch_loop(c, other_addr, no_ack=no_ack, timeout=timeout)

    c.close()

    return (internal_port, other_ip, other_port)

if __name__ == "__main__":
    parser = ap.ArgumentParser(
        description="Perform UDP holepunching using an intermediate webserver. The program returns 0 if successful and 1 otherwise; if successful, the output will have the form `<local UDP port> <remote IP> <remote UDP port>`."
    )

    subparsers = parser.add_subparsers(dest="type", required=True)

    subparser_client = subparsers.add_parser(
        "client",
        help="Serve as the client initiating the connection; requires SSH and root access to the intermediate server.",
    )
    subparser_server = subparsers.add_parser(
        "server",
        help="Serve as the server polling the intermediate server for a connection.",
    )

    for p in (subparser_client, subparser_server):
        p.add_argument(
            "-l", "--local-port", 
            type=int, 
            default=0, 
            help="Fixed internal UDP port for us to use (defaults to 0 = auto-assigned by system)"
        )
        p.add_argument(
            "-p", "--server-port", 
            type=int, 
            default=0,
            help="Fixed UDP port to use to start holepunching (defaults to 0 to use our external UDP port number)"
        )
        p.add_argument(
            "-s", "--skip-ack",
            action="store_true",
            default=False,
            help="Skip hello packet acknowledgement to finalize holepunching (a couple of the first packets down the line may get lost)",
        )

    subparser_client.add_argument(
        "server_name", 
        help="Name of the intermediate server (for SSH)"
    )
    subparser_client.add_argument(
        "server_path", 
        help="Path of file to write on the intermediate server (e.g. /var/www/inter)"
    )
    subparser_client.add_argument(
        "-k", "--signing-key",
        type=str,
        default=None,
        help="Path to the SSH key used to sign postings on the webserver"
    )

    subparser_server.add_argument(
        "server_path", 
        help="Path to check (e.g. https://example.com/something)"
    )
    subparser_server.add_argument(
        "-k", "--allowed-keys",
        type=str,
        default=None,
        help="Comma-separated list of allowed SSH public keys to verify postings on the webserver, as stored in e.g. authorized_keys"
    )

    args = parser.parse_args()

    if args.type == "client":
        ret = udp_holepunch_client(
            args.local_port,
            args.server_name,
            args.server_port,
            args.server_path,
            no_ack=args.skip_ack,
            key_path=args.signing_key,
        )
    else:
        allowed_keys = None
        if args.allowed_keys is not None:
            allowed_keys = args.allowed_keys.split(',')
        ret = udp_holepunch_server(
            args.local_port,
            args.server_port,
            args.server_path,
            no_ack=args.skip_ack,
            allowed_keys=allowed_keys,
        )
        if ret is None:
            sys.exit(1)

    print(' '.join((str(i) for i in ret)))