478 lines
14 KiB
Python
478 lines
14 KiB
Python
#!/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)
|