#!/usr/bin/env python3

import argparse as ap
import subprocess as subp
import sys

from udp_holepunch import udp_holepunch_client

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

    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(
        "-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(
        "-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(
        "-e", "--sctp-echo",
        default="sctp_echo",
        help="Path to invoke sctp_echo"
    )
    parser.add_argument(
        "-k", "--signing-key",
        type=str,
        default=None,
        help="Path to the SSH key used to sign postings on the webserver (can be different from the login key"
    )
    parser.add_argument(
        "--ack-timeout",
        type=int,
        default=60,
        help="Timeout when waiting for holepunching to be established",
    )

    parser.add_argument(
        "server_name", 
        help="Name of the intermediate server (for SSH)"
    )
    parser.add_argument(
        "server_path", 
        help="Path of file to write on the intermediate server (e.g. /var/www/inter)"
    )
    parser.add_argument(
        "ssh_name",
        help="Name of SSH server to connect to (%%h if in ProxyCommand)",
    )

    args = parser.parse_args()

    ## 1. Holepunch
    local_port, remote_ip, remote_port = 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,
        timeout=args.ack_timeout,
    )

    ## 2. Run the client
    client = subp.Popen(
        [
            args.sctp_echo, 
            "-c",
            "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 UDP port
            str(remote_port), ## Remote SCTP port = UDP
        ],
        stdout=sys.stdout.buffer.raw, ## Must use raw or e.g. vi will crash it
        stdin=sys.stdin.buffer.raw,
    )
    client.wait()