lib.grade_helpers

  1import base64
  2import posixpath
  3import re
  4import shlex
  5from ipaddress import IPv4Address, IPv4Interface, IPv6Address, IPv6Interface
  6from typing import Union
  7
  8from SRE.lib_sre import Grade0
  9
 10
 11def test_dig(grade: Grade0, machine_name: str,
 12             server_ip: Union[str, IPv4Address, IPv4Interface, IPv6Address, IPv6Interface], *,
 13             proto: str = "udp", port: int = 53, request: str,
 14             timeout: int = 2, step: int = 1) -> tuple[str, int]:
 15    """Run a ``dig +short`` query from a container and return ``(stdout, exit_code)``.
 16
 17    The command issued is::
 18
 19        dig +time={timeout-1} +tries=1 +short [+tcp] -p {port} @{server_ip} {request}
 20
 21    Args:
 22        grade:        the Grade0 instance — required so the command is
 23                      registered in the first pass and re-fetched in the second.
 24        machine_name: container the dig command is executed from.
 25        server_ip:    target DNS server. Accepts an IPv4/IPv6 literal string, an
 26                      ``IPv4Address`` / ``IPv6Address`` or an ``IPv4Interface`` /
 27                      ``IPv6Interface`` (the prefix is stripped). Substituted as
 28                      ``@server_ip`` (dig takes a bare IPv6 address, no brackets).
 29        proto:        ``"udp"`` (default) or ``"tcp"``; TCP adds ``+tcp``.
 30        port:         server port (default 53).
 31        request:      dig query body — everything that follows ``@server``,
 32                      e.g. ``"www.example.com A"`` or ``"example.com SOA"``.
 33        timeout:      total timeout in seconds passed to ``grade.test()`` (default 2).
 34                      ``dig`` itself is given ``+time={timeout-1}`` so it gives up
 35                      one second before the outer test, leaving room for clean exit.
 36        step:         step number passed to ``grade.test()`` (default 1).
 37
 38    Returns:
 39        ``(stdout, exit_code)`` — stdout is stripped. On time-outs or unreachable
 40        servers the output may be empty; the exit code from ``grade.test()`` is
 41        passed through unchanged. Errors are not recorded (``allow_error=True``).
 42    """
 43    if isinstance(server_ip, (IPv4Interface, IPv6Interface)):
 44        server_ip_str = str(server_ip.ip)
 45    elif isinstance(server_ip, (IPv4Address, IPv6Address)):
 46        server_ip_str = str(server_ip)
 47    else:
 48        server_ip_str = server_ip
 49    proto_flag = " +tcp" if proto.lower() == "tcp" else ""
 50    cmd = (f"dig +time={timeout - 1} +tries=1 +short{proto_flag} "
 51           f"-p {port} @{server_ip_str} {request}")
 52    out, code = grade.test(machine_name, cmd,
 53                           step=step, timeout=timeout, allow_error=True)
 54    return (out or "").strip(), code
 55
 56
 57def eval_tcp_server(grade: Grade0, machine_name: str, server_name: str,
 58                    step: int = 1) -> list[int] | None:
 59    """Check that a process matching server_name is running on machine_name and
 60    return the TCP ports it listens on.
 61
 62    Args:
 63        grade:        the Grade0 instance.
 64        machine_name: name of the virtual machine to inspect.
 65        server_name:  substring to match against running process names (ps output).
 66        step:         step number passed to grade.test() (default: 1).
 67
 68    Returns:
 69        A list of TCP port numbers in LISTEN state used by the process,
 70        or None if no matching process is found.
 71    """
 72    # All grade.test() calls must be made unconditionally so they are registered
 73    # in the first (registration) pass and carry real results in the second pass.
 74    pids_out, pids_code = grade.test(
 75        machine_name=machine_name,
 76        command=f"pgrep -f {server_name}",
 77        step=step,
 78    )
 79    ss_out, _ = grade.test(
 80        machine_name=machine_name,
 81        command=f"ss -tlnp",
 82        step=step,
 83    )
 84
 85    if pids_code != 0:
 86        return None
 87
 88    pids = set(pids_out.split())
 89    if not pids:
 90        return None
 91
 92    ports = []
 93    for line in ss_out.splitlines():
 94        # ss -tlnp output has a "users:(("name",pid=NNN,...))" field
 95        if "LISTEN" not in line:
 96            continue
 97        if not any(f"pid={pid}" in line for pid in pids):
 98            continue
 99        # Extract port from the local address column: 0.0.0.0:443, 127.0.0.1:53,
100        # 127.0.0.53%lo:53, *:443 (dual-stack socket), [::]:443, [::1]:9000, [fd00::1%eth0]:80
101        addr_m = re.search(r'\s+(?:\*|\[[^\]]*\](?:%\w+)?|[0-9a-fA-F.:]+(?:%\w+)?):(\d+)\s+', line)
102        if addr_m:
103            ports.append(int(addr_m.group(1)))
104
105    return sorted(set(ports))
106
107
108def transplant_files(grade: Grade0, src: str, dst: str, files: dict[str, str],
109                     download_step: int = 1, apply_step: int = 2,
110                     workdir: str = "/tmp/sre_transplant") -> dict[str, str]:
111    """Copy text files from the container *src* to the container *dst* during an evaluation.
112
113    *files* maps a path on *src* to a path on *dst* (normally under *workdir*). Each file is
114    read on *src* at *download_step* (``base64 -w0``) and, from the grade pass that follows
115    that step, re-created on *dst* at *apply_step* (``base64 -d``, mode 0600) — the two-step
116    contract of ``firewall.transplant_ruleset``: nothing is registered on *dst* for an empty
117    or missing file, so the command keys stay stable and no phantom "write of the empty
118    string" exists. *workdir* is wiped and re-created on *dst* at *download_step*, so a file
119    of a previous evaluation never survives.
120
121    Returns ``{src_path: content}`` with the decoded text of each file (``''`` until the
122    download step has run, or when the file is missing on *src*).
123    """
124    grade.test(dst, f"rm -rf {shlex.quote(workdir)}; mkdir -p {shlex.quote(workdir)}",
125               step=download_step, allow_error=True)
126    contents: dict[str, str] = {}
127    for src_path, dst_path in files.items():
128        b64, code = grade.test(src, f"base64 -w0 {shlex.quote(src_path)} 2>/dev/null",
129                               step=download_step, allow_error=True)
130        b64 = (b64 or "").strip()
131        if code != 0 or not b64:
132            contents[src_path] = ""
133            continue
134        try:
135            contents[src_path] = base64.b64decode(b64, validate=True).decode("utf-8", errors="replace")
136        except (ValueError, UnicodeDecodeError):
137            contents[src_path] = ""
138            continue
139        parent = posixpath.dirname(dst_path) or "."
140        grade.test(dst, f"mkdir -p {shlex.quote(parent)} && echo {b64} | base64 -d > {shlex.quote(dst_path)}"
141                        f" && chmod 600 {shlex.quote(dst_path)}",
142                   step=apply_step, allow_error=True)
143    return contents