#!/usr/bin/env python3

"""
The VLAB client script.

Example usage:
        ./vlab.py -k keyfile.vlabkey -b an_fpga_board

The script will assume your username is the same as the system user. To specify a different username
use the -u option.

The VLAB server and port are given as defaults to the -r and -p arguments so the user should not
have to specify them, but they can be overridden if required.

Ian Gray, 2017-2026
"""

import argparse
import os
import shlex
import signal
import socket
import sys
import time
import urllib.error
import urllib.request
from subprocess import DEVNULL, PIPE, Popen
from typing import NoReturn

############################
# Update version string here and in 'current_version' file when updating this script
# Version number must be in 'x.y.z' format
current_version = "1.3.0"
current_branch = "master"
############################


def err(error_string) -> NoReturn:
    print(error_string)
    sys.exit(1)


parser = argparse.ArgumentParser(description="VLAB client script")
parser.add_argument(
    "-r",
    "--relay",
    nargs=1,
    default=["rts001.cs.york.ac.uk"],
    help="The hostname of the relay server.",
)
parser.add_argument(
    "-p",
    "--port",
    nargs=1,
    default=["2222"],
    help="The ssh port of the relay server to connect to.",
)
parser.add_argument(
    "-l",
    "--localport",
    nargs=1,
    default=["12345"],
    help="Local port to forward connections to.",
)
parser.add_argument(
    "-k", "--key", nargs=1, help="VLAB keyfile to use for authentication."
)
parser.add_argument("-u", "--user", nargs=1, help="VLAB username.")
parser.add_argument(
    "-b", "--board", nargs=1, default=["vlab_zybo-z7"], help="Requested board class."
)
parser.add_argument("-s", "--serial", nargs=1, help="Requested board serial number.")
parser.add_argument(
    "-v",
    "--verbose",
    default=False,
    action="store_true",
    help="Enable verbose logging.",
)
parser.add_argument(
    "-c",
    "--capture",
    default=False,
    action="store_true",
    help="Capture the FPGA framebuffer and save as JPEG.",
)
parser.add_argument(
    "-o",
    "--output",
    nargs=1,
    help="Output filename for capture (default: capture_YYYYMMDD_HHMMSS.jpg)",
)
parser.add_argument(
    "--vdma",
    nargs=1,
    help="Base address of the framebuffer VDMA peripheral for --capture, e.g. 0x43000000 "
    "(default: 0x43000000).",
)
parser.add_argument(
    "-n",
    "--no-terminal",
    default=False,
    action="store_true",
    help="Hold the board and the port forward open without starting a terminal. "
    "Leaves the board's serial port free for --capture-serial. For scripting.",
)
parser.add_argument(
    "--capture-serial",
    nargs=1,
    metavar="SECONDS",
    help="Capture the board's serial output for the given number of seconds and write "
    "it to stdout, or to --output if given. Needs a session started with "
    "--no-terminal.",
)
parsed = parser.parse_args()

error_info = (
    "Read the instructions at\n"
    "\thttps://iangray001.github.io/embs/docs/setup/vlab-quickstart/"
)


# Check for an update from GitHub repository
update_url = "https://raw.githubusercontent.com/RTSYork/VLAB/{}/client_version".format(
    current_branch
)

if parsed.verbose:
    print("Checking for update at URL: {}".format(update_url))

req = urllib.request.Request(update_url)
try:
    response = urllib.request.urlopen(req)
    content = response.readline()
    remote_version_string = content.decode("utf-8")
    remote_version = tuple(map(int, (remote_version_string.split("."))))
    local_version = tuple(map(int, (current_version.split("."))))
    if local_version < remote_version:
        print("A new version of the VLAB script is available on GitHub")
        print(
            "You have version v{} and the latest is v{}".format(
                current_version, remote_version_string
            )
        )
        print(
            "Download the latest version from https://raw.githubusercontent.com/RTSYork/VLAB/{}/vlab.py".format(
                current_branch
            )
        )
        print()
        try:
            input("Press Control-C to exit, or press Enter to continue...")
        except KeyboardInterrupt:
            sys.exit(0)
except urllib.error.URLError as e:
    print("Error checking for script updates:", e.reason)

# Check the keyfile exists
if not parsed.key:
    err("Specify a keyfile with --key.")
if not os.path.isfile(parsed.key[0]):
    err(
        "Keyfile {} does not exist. Specify a keyfile with --key.".format(parsed.key[0])
    )

# Handle capture mode
if parsed.capture:
    ssh_cmd = [
        "ssh",
        "-oPasswordAuthentication=no",
        "-i",
        parsed.key[0],
        "-p",
        parsed.port[0],
    ]
    if parsed.user is not None:
        ssh_cmd.extend(["-l", parsed.user[0]])
    # Encode an optional VDMA base address into the relay command token (capture[:vdma]),
    # mirroring the boardclass:port:serial form used for board sessions. The relay validates
    # the value before use; omitting it lets the boardserver fall back to its built-in default.
    if parsed.vdma is not None:
        ssh_cmd.extend([parsed.relay[0], "capture:{}".format(parsed.vdma[0])])
    else:
        ssh_cmd.extend([parsed.relay[0], "capture"])

    if parsed.verbose:
        print("Capture ssh command: {}".format(ssh_cmd))

    print("Requesting framebuffer capture...")
    proc = Popen(ssh_cmd, stdout=PIPE)
    stdout, _ = proc.communicate()

    if proc.returncode != 0 or len(stdout) == 0:
        if len(stdout) > 0:
            print(stdout.decode("utf-8", errors="replace").strip())
        else:
            print("Capture failed.")
        sys.exit(1)

    # Verify JPEG magic bytes
    if len(stdout) < 2 or stdout[:2] != b"\xff\xd8":
        print(stdout.decode("utf-8", errors="replace").strip())
        sys.exit(1)

    if parsed.output:
        filename = parsed.output[0]
    else:
        filename = "capture_{}.jpg".format(time.strftime("%Y%m%d_%H%M%S"))

    with open(filename, "wb") as f:
        f.write(stdout)
    print("Saved to {} ({} bytes)".format(filename, len(stdout)))
    sys.exit(0)

# Handle serial capture mode. Like --capture this is a one-shot command against a
# session the user already holds, so it needs no port forward of its own.
if parsed.capture_serial:
    seconds = parsed.capture_serial[0]
    if not seconds.isdigit() or int(seconds) < 1:
        err("--capture-serial takes a whole number of seconds.")

    ssh_cmd = [
        "ssh",
        "-oPasswordAuthentication=no",
        "-i",
        parsed.key[0],
        "-p",
        parsed.port[0],
    ]
    if parsed.user is not None:
        ssh_cmd.extend(["-l", parsed.user[0]])
    ssh_cmd.extend([parsed.relay[0], "serial:{}".format(seconds)])

    if parsed.verbose:
        print("Serial capture ssh command: {}".format(ssh_cmd), file=sys.stderr)

    # Only the captured bytes go to stdout, so that the caller can redirect them.
    # The relay writes its own progress and errors to stderr, which passes through.
    proc = Popen(ssh_cmd, stdout=PIPE)
    captured, _ = proc.communicate()

    if proc.returncode != 0:
        sys.exit(proc.returncode)

    if parsed.output:
        with open(parsed.output[0], "wb") as f:
            f.write(captured)
        print(
            "Saved {} bytes to {}".format(len(captured), parsed.output[0]),
            file=sys.stderr,
        )
    else:
        sys.stdout.buffer.write(captured)
        sys.stdout.buffer.flush()
    sys.exit(0)

# Check that the requested ports are free to use
local_port = int(parsed.localport[0])
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
    local_port_in_use = s.connect_ex(("localhost", local_port)) == 0

if local_port_in_use:
    err(
        "Error: Local port {} is in use by another user on this machine, or otherwise unavailable. "
        "Specify a different port with --localport or -l.\n"
        "\nEach user must choose their own unique local port when running on a shared machine (e.g. a departmental "
        "server).\nSee the section on 'Using the VLAB on a Shared Machine' at "
        "https://wiki.york.ac.uk/display/RTS/The+VLAB+Quickstart+Guide for more information.".format(
            local_port
        )
    )

# First get a port to use on the relay server
if parsed.user is not None:
    ssh_cmd = [
        "ssh",
        "-oPasswordAuthentication=no",
        "-i",
        parsed.key[0],
        "-p",
        parsed.port[0],
        "-l",
        parsed.user[0],
        parsed.relay[0],
        "getport",
    ]
else:
    ssh_cmd = [
        "ssh",
        "-oPasswordAuthentication=no",
        "-i",
        parsed.key[0],
        "-p",
        parsed.port[0],
        parsed.relay[0],
        "getport",
    ]
if parsed.verbose:
    print("First ssh command: {}".format(ssh_cmd))
stdout, stderr = Popen(ssh_cmd, stdout=PIPE, stderr=PIPE).communicate()

try:
    error = stderr.decode("ascii").strip()
except UnicodeError:
    error = ""
    reply = None
    err("Could not decode error output from ssh subprocess")

if len(stdout) == 0 and len(error) > 0:
    errstr = ""
    if error.find("UNPROTECTED PRIVATE KEY FILE") != -1:
        errstr += "The permissions on the key {} are too permissive.\n".format(
            parsed.key[0]
        )
        errstr += "Only your user should be able to access the key, or it will be rejected by ssh.\n"
        errstr += "To fix this, on Linux run the command\n\tchmod 0600 {}\n".format(
            parsed.key[0]
        )
        errstr += "On Windows, see the instructions at the link below.\n"
    elif (
        error.find("Operation timed out") != -1
        or error.find("Connection refused") != -1
    ):
        errstr += "Cannot see the VLAB relay server: {}\n".format(parsed.relay[0])
        errstr += "Try the following:\n\tssh {}\n".format(parsed.relay[0])
        errstr += "You should get a password prompt:\n\t> <yourusername>@{}'s password\n".format(
            parsed.relay[0]
        )
        errstr += "If this doesn't connect, or the prompt is for any other server, set up your .ssh/config file\n"
    elif error.find("Permission denied") != -1:
        errstr += error
        errstr += "\nEnsure that you have copied your private key:\n"
        errstr += "\tssh-copy-id csteach1.york.ac.uk\n\t(or csresearch1.york.ac.uk for staff)\n"
        errstr += "Also, if your VLAB username is different to your normal username, set it with -u:\n"
        errstr += "\t./vlab.py -k {} -u myusername\n".format(parsed.key[0])
    elif error.find("Resource temporarily unavailable") != -1:
        errstr += error
        errstr += (
            "\nThis often means a firewall has blocked the connection. Try temporarily disabling your firewall "
            "application.\n"
        )
    elif error.find("Connection closed by") != -1:
        errstr += error
        errstr += (
            "\nEnsure that you have set the correct usernames in your ssh config.\n"
        )
    elif error.find("REMOTE HOST IDENTIFICATION HAS CHANGED") != -1:
        errstr += error
        errstr += (
            "\n\n\nThe host key in the VLAB has changed. This should not happen and might indicate a "
            "man-in-the-middle attack.\n"
        )
        errstr += "If you are sure this is not the case, delete the line mentioned above from the file ~/.ssh/known_hosts\n"
    else:
        print(error)
    errstr += error_info
    err(errstr)

try:
    reply = stdout.decode("ascii").strip()
except UnicodeError:
    reply = None
    err(
        "Invalid reply from VLAB server. Incorrect message format.\n{}".format(
            error_info
        )
    )

if len(reply) > 9:
    if reply[:9] == "VLABPORT:":
        port_string = reply[9:]
        ephemeral_port = int(port_string)
    else:
        ephemeral_port = None
        err("Invalid reply from VLAB server. Wrong header.\n{}".format(error_info))
else:
    ephemeral_port = None
    err("Invalid reply from VLAB server. Message too short.\n{}".format(error_info))

print(
    "Tunnelling to relay server '{}' using port {}...".format(
        parsed.relay[0], ephemeral_port
    )
)

# Now create the actual connection
ssh_user = ""
if parsed.user is not None:
    ssh_user = " -l {} ".format(parsed.user[0])

# Assemble the command we send to the relay shell
if parsed.serial is not None:
    relay_command = "{}:{}:{}".format(parsed.board[0], ephemeral_port, parsed.serial[0])
else:
    relay_command = "{}:{}".format(parsed.board[0], ephemeral_port)

# Ask the relay to allocate and tunnel as usual but not to start a terminal, so
# that the board's serial port stays free for --capture-serial.
if parsed.no_terminal:
    relay_command = "noterm:{}".format(relay_command)


ssh_cmd = (
    "ssh -L {}:localhost:{} -o PasswordAuthentication=no -o ExitOnForwardFailure=yes "
    "-e none -i {} {} -p {} -tt {} {}".format(
        local_port,
        ephemeral_port,
        parsed.key[0],
        ssh_user,
        parsed.port[0],
        parsed.relay[0],
        relay_command,
    )
)

if parsed.verbose:
    print("Second ssh command: {}".format(ssh_cmd))

if not parsed.no_terminal:
    os.system(ssh_cmd)
    sys.exit(0)

# Not shell=True: the board is released when ssh disconnects, so this process must
# be able to signal ssh itself rather than an intervening /bin/sh.
proc = Popen(
    shlex.split(ssh_cmd), stdin=DEVNULL, stdout=PIPE, bufsize=1, universal_newlines=True
)

if proc.stdout is None:
    err("Failed to capture stdout from SSH subprocess.")


def end_session(signum=None, frame=None):
    _ = (signum, frame)  # discard the parameters to avoid unused error
    if proc.poll() is None:
        proc.terminate()


# A script will stop this with SIGTERM, and a closing terminal sends SIGHUP.
# Either way the board must be handed back rather than left locked.
for _sig in (signal.SIGTERM, signal.SIGHUP):
    signal.signal(_sig, end_session)

ready = False
try:
    for line in proc.stdout:
        line = line.rstrip()
        if not ready and line.endswith("VLAB_SESSION_READY"):
            ready = True
            print("VLAB_READY:{}".format(local_port), flush=True)
            print(
                "Board ready. hw_server is forwarded to localhost:{}.".format(
                    local_port
                ),
                flush=True,
            )
            print(
                "Leave this running while you use the board; the lock is released when it exits.",
                flush=True,
            )
            continue
        if line:
            print(line, flush=True)
except KeyboardInterrupt:
    pass
finally:
    end_session()
    try:
        proc.wait(timeout=30)
    except Exception:
        proc.kill()
        proc.wait()

if not ready:
    err("The board session never became ready.")

sys.exit(proc.returncode)