#!/usr/bin/env python3

import argparse as ap
import subprocess as subp
import sys

from udp_holepunch import udp_holepunch_server

if __name__ == "__main__":
    parser = ap.ArgumentParser(
        description="Run SSH through UDP with holepunching (server)"
    )

    parser.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)"
    )
    parser.add_argument(
        "--local-ip",
        default="0.0.0.0", 
        help="Local IP to listen on (defaults to 0.0.0.0 = all)",
    )
    parser.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)"
    )
    parser.add_argument(
        "-t", "--timeout",
        type=int, 
        default=60,
        help="Inactivity timeout once holepunching has been established (defaults to 60)",
    )
    parser.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)",
    )
    parser.add_argument(
        "--ack-timeout",
        type=int,
        default=60,
        help="Timeout when waiting for holepunching to be established",
    )
    parser.add_argument(
        "-e", "--sctp-echo",
        default="sctp_echo",
        help="Path to invoke sctp_echo"
    )
    parser.add_argument(
        "--ssh-port",
        type=int,
        default=22,
        help="SSH port on the local machine (defaults to 22)",
    )
    parser.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"
    )

    parser.add_argument(
        "server_path", 
        help="Path of file to write on the intermediate server (e.g. /var/www/inter)"
    )

    args = parser.parse_args()

    allowed_keys = None
    if args.allowed_keys is not None:
        allowed_keys = args.allowed_keys.split(',')

    ## 1. Holepunch
    ret = udp_holepunch_server(
        args.local_port,
        args.server_port,
        args.server_path,
        no_ack=args.skip_ack,
        allowed_keys=allowed_keys,
        timeout=args.ack_timeout,
    )

    if ret is None:
        ## Nothing on the intermediate server, just exit
        sys.exit(0)

    local_port, remote_ip, remote_port = ret

    ## 2. Run the server
    server = subp.Popen(
        [
            args.sctp_echo, 
            "-s",
            "0.0.0.0", ## Local IP
            str(local_port), ## Local UDP
            str(local_port), ## Local SCTP = UDP
            remote_ip, ## Remote IP
            str(remote_port), ## Remote SCTP port
        ],
        stdout=subp.PIPE,
        stdin=subp.PIPE,
    )

    ## 3. Use socat to make the actual connection
    socat = subp.Popen(
        [
            "socat", 
            "-T", str(args.timeout),
            "-", 
            "TCP4:127.0.0.1:%d" % args.ssh_port,
        ], 
        stdin=server.stdout,
        stdout=server.stdin,
    )

    ## Communicate with socat; this lets us bypass the fact that sctp-echo doesn't
    ## have timeout functionality
    socat.wait()
    server.terminate()