#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
"""Fake `ssh` for `ameesh offload` tests (L159b): never a network.

`AMEESH_OFFLOAD_SSH=tests/fakebin/fake-offload-ssh` makes the client call it
like ssh (`[-o OPT]... HOST COMMAND`). It runs the VM side,
`ameesh.platform.offload_shell --config $AMEESH_FAKE_OFFLOAD_CONFIG`, with the command
in SSH_ORIGINAL_COMMAND, as sshd does for a forced command. The host
`unreachable` exits 255, as ssh does when it cannot connect.
"""
import os
import sys

args = sys.argv[1:]
while args and args[0] in ("-o", "-p", "-i", "-F"):
    args = args[2:]
if len(args) != 2:
    sys.stderr.write("fake-offload-ssh: expected HOST COMMAND, got %r\n" % args)
    sys.exit(255)
host, command = args
if host == "unreachable":
    sys.stderr.write("ssh: connect to host unreachable port 22: Connection refused\n")
    sys.exit(255)
env = dict(os.environ, SSH_ORIGINAL_COMMAND=command)
src = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "src")
env["PYTHONPATH"] = src
# a child, not exec: the shell sees its SSH "session" (this process) die.
# Its stdin, stdout and stderr go through this process, as through sshd (the
# shell holds no end of the client's own pipes): killing it closes the
# client's pipes (the "network" drops) while the shell lives on.
import subprocess
import threading

proc = subprocess.Popen([sys.executable, "-m", "ameesh.platform.offload_shell",
                         "--config", os.environ["AMEESH_FAKE_OFFLOAD_CONFIG"]], env=env,
                        stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE)


def relay(src, dst, close=False):
    for chunk in iter(lambda: src.read1(65536), b""):
        try:
            dst.write(chunk)
            dst.flush()
        except OSError:
            break
    if close:
        try:
            dst.close()
        except OSError:
            pass


threading.Thread(target=relay, args=(sys.stdin.buffer, proc.stdin, True),
                 daemon=True).start()
threads = [threading.Thread(target=relay, args=(proc.stdout, sys.stdout.buffer)),
           threading.Thread(target=relay, args=(proc.stderr, sys.stderr.buffer))]
for thread in threads:
    thread.start()
try:
    rc = proc.wait()
except KeyboardInterrupt:
    rc = proc.wait()
for thread in threads:
    thread.join()
sys.stdout.flush()
sys.stderr.flush()
# _exit: the stdin relay may still be blocked in a read, as ssh would be
os._exit(rc if rc >= 0 else 128 - rc)
