#! /somewhere/python3
import argparse
import atexit
import base64
import io
import logging
import os
import pathlib
import pty
import queue
import re
import shlex
import shutil
import socket
import subprocess
import sys
import tempfile
import _thread
import time
from contextlib import contextmanager
from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
import ptpython.repl
CHAR_TO_KEY = {
"A": "shift-a",
"N": "shift-n",
"-": "0x0C",
"_": "shift-0x0C",
"B": "shift-b",
"O": "shift-o",
"=": "0x0D",
"+": "shift-0x0D",
"C": "shift-c",
"P": "shift-p",
"[": "0x1A",
"{": "shift-0x1A",
"D": "shift-d",
"Q": "shift-q",
"]": "0x1B",
"}": "shift-0x1B",
"E": "shift-e",
"R": "shift-r",
";": "0x27",
":": "shift-0x27",
"F": "shift-f",
"S": "shift-s",
"'": "0x28",
'"': "shift-0x28",
"G": "shift-g",
"T": "shift-t",
"`": "0x29",
"~": "shift-0x29",
"H": "shift-h",
"U": "shift-u",
"\\": "0x2B",
"|": "shift-0x2B",
"I": "shift-i",
"V": "shift-v",
",": "0x33",
"<": "shift-0x33",
"J": "shift-j",
"W": "shift-w",
".": "0x34",
">": "shift-0x34",
"K": "shift-k",
"X": "shift-x",
"/": "0x35",
"?": "shift-0x35",
"L": "shift-l",
"Y": "shift-y",
" ": "spc",
"M": "shift-m",
"Z": "shift-z",
"\n": "ret",
"!": "shift-0x02",
"@": "shift-0x03",
"#": "shift-0x04",
"$": "shift-0x05",
"%": "shift-0x06",
"^": "shift-0x07",
"&": "shift-0x08",
"*": "shift-0x09",
"(": "shift-0x0A",
")": "shift-0x0B",
}
# Forward reference
machines: "List[Machine]"
logging.basicConfig(format="%(message)s")
logger = logging.getLogger("test-driver")
logger.setLevel(logging.INFO)
class MachineLogAdapter(logging.LoggerAdapter):
def process(self, msg: str, kwargs: Any) -> Tuple[str, Any]:
return f"{self.extra['machine']}: {msg}", kwargs
def make_command(args: list) -> str:
return " ".join(map(shlex.quote, (map(str, args))))
def create_vlan(vlan_nr: str) -> Tuple[str, str, "subprocess.Popen[bytes]", Any]:
logger.info(f"starting VDE switch for network {vlan_nr}")
vde_socket = tempfile.mkdtemp(
prefix="nixos-test-vde-", suffix="-vde{}.ctl".format(vlan_nr)
)
pty_master, pty_slave = pty.openpty()
vde_process = subprocess.Popen(
["vde_switch", "-s", vde_socket, "--dirmode", "0700"],
bufsize=1,
stdin=pty_slave,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
shell=False,
)
fd = os.fdopen(pty_master, "w")
fd.write("version\n")
# TODO: perl version checks if this can be read from
# an if not, dies. we could hang here forever. Fix it.
assert vde_process.stdout is not None
vde_process.stdout.readline()
if not os.path.exists(os.path.join(vde_socket, "ctl")):
raise Exception("cannot start vde_switch")
return (vlan_nr, vde_socket, vde_process, fd)
def retry(fn: Callable) -> None:
"""Call the given function repeatedly, with 1 second intervals,
until it returns True or a timeout is reached.
"""
for _ in range(900):
if fn(False):
return
time.sleep(1)
if not fn(True):
raise Exception("action timed out")
class Machine:
def __init__(self, args: Dict[str, Any]) -> None:
if "name" in args:
self.name = args["name"]
else:
self.name = "machine"
cmd = args.get("startCommand", None)
if cmd:
match = re.search("run-(.+)-vm$", cmd)
if match:
self.name = match.group(1)
self.script = args.get("startCommand", self.create_startcommand(args))
tmp_dir = os.environ.get("TMPDIR", tempfile.gettempdir())
def create_dir(name: str) -> str:
path = os.path.join(tmp_dir, name)
os.makedirs(path, mode=0o700, exist_ok=True)
return path
self.state_dir = create_dir("vm-state-{}".format(