#!/usr/bin/env python3
"""A stand-in for the `ssh` binary, for tests that must not need an SSH server.

It understands the four things Harlequin passes it -- `-G`, `-V`, `-L` and a
destination -- and answers `-G` in the format OpenSSH's own config dumper uses.
Without `-G` it binds the local end of each forward and sleeps, which is every
observable thing a real forward does to the machine Harlequin is running on.

Environment variables shape what it does, so one script covers the cases:

    FAKE_SSH_ARGV       append each invocation's argv, as a JSON line, to this path
    FAKE_SSH_FORWARD    a `LocalForward` the destination's Host block supplies
    FAKE_SSH_STDERR     write this to stderr, as ssh writes its diagnostics
    FAKE_SSH_STDERR_PARTIAL  write this to stderr with no trailing newline
    FAKE_SSH_STDOUT     write this to stdout, as a ProxyCommand helper may
    FAKE_SSH_EXIT       exit with this code instead of binding anything
    FAKE_SSH_HANG       bind nothing and sleep, as ssh does waiting on a passphrase
    FAKE_SSH_DELAY      seconds to wait before binding
    FAKE_SSH_DROP_WHEN  hold the forwards until this path exists, then drop them
    FAKE_SSH_PROBE_EXIT exit `-G` with this code
    FAKE_SSH_PROBE_TEXT print this from `-G` instead of a resolved config
    FAKE_SSH_ALIVE_INTERVAL  what `-G` reports for serveraliveinterval (default 0,
                        so every tunnel built here is given a keepalive)
"""

from __future__ import annotations

import json
import os
import socket
import sys
import time

TAKES_A_VALUE = {"-L", "-R", "-D", "-o", "-i", "-p", "-l", "-F", "-b", "-c", "-W"}


def parse_argv(argv: list[str]) -> tuple[list[str], str, bool]:
    """The `-L` specs, the destination, and whether `-G` was asked for."""
    forwards: list[str] = []
    destination = ""
    dump_config = False
    index = 0
    while index < len(argv):
        arg = argv[index]
        if arg in TAKES_A_VALUE:
            value = argv[index + 1] if index + 1 < len(argv) else ""
            if arg == "-L":
                forwards.append(value)
            index += 2
        elif arg.startswith("-"):
            dump_config = dump_config or "G" in arg
            index += 1
        else:
            if not destination:
                destination = arg
            index += 1
    return forwards, destination, dump_config


def format_forward(spec: str) -> tuple[str, str]:
    """One `-L` spec as the two fields `ssh -G` prints for it."""
    parts = spec.split(":")
    if len(parts) == 4:
        bind, listen_port, host, port = parts
        listen = f"[{bind}]:{listen_port}"
    else:
        listen_port, host, port = parts
        listen = listen_port
    return listen, f"[{host}]:{port}"


def endpoint(listen: str) -> tuple[str, int]:
    """Where the local end of a forward listens."""
    if listen.startswith("["):
        host, _, port = listen[1:].partition("]:")
        return host, int(port)
    return "127.0.0.1", int(listen)


def main() -> int:
    argv = sys.argv[1:]
    record = os.environ.get("FAKE_SSH_ARGV")
    if record:
        with open(record, "a", encoding="utf-8") as f:
            f.write(json.dumps(argv) + "\n")

    if "-V" in argv:
        print("FakeSSH_1.0, a stand-in", file=sys.stderr)
        return 0

    forwards, destination, dump_config = parse_argv(argv)
    from_config = os.environ.get("FAKE_SSH_FORWARD")
    if from_config:
        forwards.append(from_config)

    if dump_config:
        probe_exit = int(os.environ.get("FAKE_SSH_PROBE_EXIT", "0"))
        if probe_exit:
            # a real client says why it would not take the argv
            message = os.environ.get("FAKE_SSH_STDERR")
            if message:
                print(message, file=sys.stderr, flush=True)
            return probe_exit
        text = os.environ.get("FAKE_SSH_PROBE_TEXT")
        if text is not None:
            sys.stdout.write(text)
            return 0
        print(f"host {destination}")
        print(f"hostname {destination}")
        print("port 22")
        print(f"serveraliveinterval {os.environ.get('FAKE_SSH_ALIVE_INTERVAL', '0')}")
        for spec in forwards:
            listen, connect = format_forward(spec)
            print(f"localforward {listen} {connect}")
        return 0

    message = os.environ.get("FAKE_SSH_STDERR")
    if message:
        print(message, file=sys.stderr, flush=True)
    on_stdout = os.environ.get("FAKE_SSH_STDOUT")
    if on_stdout:
        print(on_stdout, flush=True)
    partial = os.environ.get("FAKE_SSH_STDERR_PARTIAL")
    if partial:
        # no trailing newline, as a prompt has none
        sys.stderr.write(partial)
        sys.stderr.flush()
    code = os.environ.get("FAKE_SSH_EXIT")
    if code:
        return int(code)

    time.sleep(float(os.environ.get("FAKE_SSH_DELAY", "0")))
    listeners = []
    if not os.environ.get("FAKE_SSH_HANG"):
        for spec in forwards:
            listen, _ = format_forward(spec)
            host, port = endpoint(listen)
            sock = socket.socket()
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            try:
                sock.bind((host, port))
            except OSError as e:
                print(f"bind [{host}]:{port}: {e.strerror}", file=sys.stderr)
                return 255
            sock.listen(8)
            listeners.append(sock)
    drop_when = os.environ.get("FAKE_SSH_DROP_WHEN")
    if drop_when:
        # the test controls when the forward drops, so nothing races the
        # readiness poll on a loaded machine
        while not os.path.exists(drop_when):
            time.sleep(0.02)
        for sock in listeners:
            sock.close()
        print("Timeout, server web-1 not responding.", file=sys.stderr)
        return 255
    while True:
        time.sleep(3600)


if __name__ == "__main__":
    sys.exit(main())
