diff --git a/src/mcbebop/files/ftp.py b/src/mcbebop/files/ftp.py new file mode 100644 index 0000000..6b4d045 --- /dev/null +++ b/src/mcbebop/files/ftp.py @@ -0,0 +1,265 @@ +"""The drone's FTP servers, read-only. + +Firmware 4.7.1 runs three separate `ftpd` instances, all as root, all +anonymous, and all confirmed open on the live aircraft with no button press +needed. They differ only in which directory they are chrooted to: + +| Port | Root | What lives there | +|------|----------------------------------|-----------------------------------| +| 21 | `/data/ftp` | media, plus `internal_000/Debug/` | +| 61 | `/data/ftp/internal_000/flightplans` | `flightPlan.mavlink` | +| 51 | `/update` | firmware images | + +This module exposes the first two, split into three named areas so that an +agent asking for logs cannot wander into the media tree by accident. + +Port 51 is deliberately absent, and that absence is the point: it is the +channel firmware is *pushed* through. There is no read use case for it, the +drone's Wi-Fi is open with no passphrase, and a write there replaces the +operating system of an aircraft that flies over people. Exposing it as a tool +would hand the weakest point of the whole machine to anything that can talk to +this server. If a firmware image is ever genuinely needed, do it by hand, on +purpose, with the drone on a bench. + +Everything here is anonymous login with bounded timeouts, and fetches stream +to a file under `Settings.capture_dir` rather than into memory: the media tree +holds 1080p recordings, and pulling one into a dict that becomes a tool result +would be a bad day for whoever is paying for context. +""" + +from __future__ import annotations + +import ftplib +import logging +from dataclasses import dataclass +from pathlib import Path + +logger = logging.getLogger(__name__) + +DEFAULT_TIMEOUT = 10.0 +#: Enough for a log bundle or a flight plan, far short of a video. +DEFAULT_MAX_BYTES = 32 * 1024 * 1024 + +# Named, not wired to anything. Here so that a reader looking for "what about +# the update port?" finds the answer next to the code rather than in a commit +# message. See the module docstring. +FIRMWARE_WRITE_PORT = 51 + + +class FtpError(RuntimeError): + """The drone's FTP server refused, or the request was not allowed.""" + + +class TooLarge(FtpError): + """The file is bigger than the caller said it would accept.""" + + +@dataclass(frozen=True) +class Area: + """One read-only view of the drone's storage.""" + + name: str + port: int + root: str # the chroot the drone's ftpd serves, for error messages only + prefix: str # scopes this area to a subtree of that chroot + describe: str + + +AREAS: dict[str, Area] = { + "media": Area( + name="media", + port=21, + root="/data/ftp", + prefix="", + describe="photos and videos the drone recorded, under internal_000/", + ), + "flightplans": Area( + name="flightplans", + port=61, + root="/data/ftp/internal_000/flightplans", + prefix="", + describe="MAVLink flight plans; this aircraft ships a flightPlan.mavlink", + ), + "logs": Area( + # Same ftpd as media, narrowed to the blackbox tree. The drone writes + # its flight logs here and nothing else in this subtree is interesting. + name="logs", + port=21, + root="/data/ftp", + prefix="internal_000/Debug/current", + describe="blackbox flight logs the drone writes during each session", + ), +} + + +@dataclass(frozen=True) +class Entry: + """One line of a directory listing.""" + + name: str + size: int | None + is_dir: bool + raw: str + + +def areas() -> list[Area]: + """Every area a caller may read. Port 51 is not among them, by design.""" + return list(AREAS.values()) + + +def resolve_area(name: str) -> Area: + area = AREAS.get(name) + if area is None: + raise FtpError( + f"No FTP area named {name!r}. Available: {', '.join(sorted(AREAS))}. " + "The drone's firmware-update port is not reachable from here on purpose." + ) + return area + + +def _safe_path(area: Area, path: str) -> str: + """Join `path` into the area, refusing anything that could escape it. + + The drone's ftpd is chrooted, so this is not protecting the aircraft. It is + protecting the *area* boundary, which is the only thing keeping a request + for logs from reading the whole media tree. + """ + cleaned = path.strip().strip("/") + parts = [p for p in cleaned.split("/") if p not in ("", ".")] + if any(p == ".." for p in parts): + raise FtpError(f"'..' is not allowed in an FTP path (got {path!r}).") + if path.startswith("/"): + raise FtpError(f"FTP paths are relative to the {area.name} area; drop the leading '/' from {path!r}.") + joined = "/".join(filter(None, [area.prefix, *parts])) + return joined or "." + + +def _connect(area: Area, host: str, timeout: float) -> ftplib.FTP: + try: + ftp = ftplib.FTP(timeout=timeout) + ftp.connect(host, area.port) + ftp.login() # anonymous; the drone asks for nothing + except OSError as exc: + raise FtpError( + f"Could not reach the drone's {area.name} FTP server at {host}:{area.port} ({exc}). " + "Check this machine has joined the aircraft's access point." + ) from exc + return ftp + + +def _parse_line(line: str) -> Entry: + """Parse a unix-style LIST line, degrading to a bare name. + + busybox ftpd emits the usual `drwxr-xr-x 2 root root 4096 Jan 1 00:00 name` + shape, but there is no standard here, so anything unexpected still yields a + usable name and keeps the raw line for a human to read. + """ + fields = line.split(maxsplit=8) + if len(fields) == 9 and len(fields[0]) == 10: + size = int(fields[4]) if fields[4].isdigit() else None + return Entry(name=fields[8], size=size, is_dir=line.startswith("d"), raw=line) + return Entry(name=line.strip(), size=None, is_dir=False, raw=line) + + +def list_dir( + area_name: str, + path: str = "", + *, + host: str, + timeout: float = DEFAULT_TIMEOUT, +) -> list[Entry]: + """List one directory inside an area.""" + area = resolve_area(area_name) + target = _safe_path(area, path) + ftp = _connect(area, host, timeout) + try: + lines: list[str] = [] + ftp.retrlines(f"LIST {target}", lines.append) + except ftplib.all_errors as exc: + raise FtpError(f"Could not list {target!r} in the {area.name} area: {exc}") from exc + finally: + _close(ftp) + return [_parse_line(line) for line in lines if line.strip()] + + +def fetch( + area_name: str, + path: str, + *, + host: str, + capture_dir: Path, + max_bytes: int = DEFAULT_MAX_BYTES, + timeout: float = DEFAULT_TIMEOUT, +) -> Path: + """Download one file into `capture_dir` and return where it landed. + + The size is checked twice: once with SIZE before any bytes move, and again + while they do. The first check is the useful one, the second is what saves + us when the drone's ftpd declines to answer SIZE at all. + """ + area = resolve_area(area_name) + target = _safe_path(area, path) + if target == ".": + raise FtpError("fetch needs a file path, not a directory.") + + dest_dir = Path(capture_dir) / area.name + dest_dir.mkdir(parents=True, exist_ok=True) + dest = dest_dir / Path(target).name + + ftp = _connect(area, host, timeout) + try: + declared = _size_of(ftp, target) + if declared is not None and declared > max_bytes: + raise TooLarge( + f"{target} is {declared} bytes, over the {max_bytes}-byte cap. " + "Raise max_bytes if you really want it, or fetch a smaller file." + ) + written = 0 + with dest.open("wb") as fh: + + def sink(chunk: bytes) -> None: + nonlocal written + written += len(chunk) + if written > max_bytes: + # ftplib has no "stop here", so aborting the transfer means + # raising out of the callback and letting the socket close. + raise TooLarge(f"{target} exceeded the {max_bytes}-byte cap mid-transfer.") + fh.write(chunk) + + ftp.retrbinary(f"RETR {target}", sink) + except TooLarge: + dest.unlink(missing_ok=True) + raise + except ftplib.all_errors as exc: + dest.unlink(missing_ok=True) + raise FtpError(f"Could not fetch {target!r} from the {area.name} area: {exc}") from exc + finally: + _close(ftp) + + logger.info("fetched %s (%d bytes) to %s", target, written, dest) + return dest + + +def _size_of(ftp: ftplib.FTP, target: str) -> int | None: + """SIZE, or None if this ftpd will not say. + + busybox needs binary mode before it answers SIZE, and some builds refuse + regardless, so a failure here is informational rather than fatal. + """ + try: + ftp.voidcmd("TYPE I") + return ftp.size(target) + except ftplib.all_errors: + return None + + +def _close(ftp: ftplib.FTP) -> None: + """Quit politely, fall back to dropping the socket. + + A drone that has already gone away makes QUIT raise, and that must not be + what the caller sees instead of their actual result. + """ + try: + ftp.quit() + except ftplib.all_errors: + ftp.close() diff --git a/src/mcbebop/files/shell.py b/src/mcbebop/files/shell.py new file mode 100644 index 0000000..f0a6cf0 --- /dev/null +++ b/src/mcbebop/files/shell.py @@ -0,0 +1,267 @@ +"""Read-only inspection over the drone's debug telnet. + +Port 23 is closed on a freshly booted Bebop 2. Four short presses of the power +button run `shortpress_4.sh`, which calls `shpoison_cli --start_debug` and +brings up `telnetd -l /bin/login.sh`. `/bin/login.sh` is one line, +`exec /bin/sh -l`, so there is no password prompt and no user: the first thing +on the socket is a root shell. It does not survive a reboot. + +That is the whole reason this module is shaped the way it is. A general shell +behind a deny-list would be a remote root tool on an aircraft, reachable by +anyone in range of an open Wi-Fi network, and deny-lists on shells leak, +because there is always one more way to spell the thing you blocked. So the +only commands that exist here are the eleven in `ALLOWED`, every one of which +reads, and any argument carrying a shell metacharacter is refused before a +byte is sent. What goes down the socket is therefore always one simple command +with literal arguments: nothing an argument contains can become a pipe, a +redirection, a substitution or a second command. The only shell syntax in the +line is the `2>&1` this module adds itself, so a failure explains itself. +Read-only by construction rather than by policy. + +The cost of refusing quotes is that a path containing a space cannot be +expressed. Nothing on this aircraft has one, and the alternative is a quoting +layer, which is where this kind of check usually goes wrong. + +`telnetlib` was removed in Python 3.13, so the handful of protocol bytes we +need are implemented here. We never want an option, so negotiation is answered +with a flat refusal and the rest of the stream is plain text. +""" + +from __future__ import annotations + +import logging +import socket +import uuid +from dataclasses import dataclass + +logger = logging.getLogger(__name__) + +TELNET_PORT = 23 +DEFAULT_TIMEOUT = 8.0 +#: A `cat` of something unexpected should not be able to exhaust memory. +DEFAULT_MAX_BYTES = 256 * 1024 + +#: Every command this module can run. Each one only reads. +ALLOWED = frozenset( + { + "getprop", # Parrot's property store: serials, versions, build + "cat", # /proc, /sys, config files + "ls", + "dmesg", # kernel ring buffer, where driver gripes land + "ps", + "uptime", + "bcmwl", # Broadcom wireless tool: regulatory domain, channel, rates + "df", + "mount", + "head", + "tail", + } +) + +# Anything that would let an argument become more than an argument. Backslash +# is in here because it is how you smuggle the rest of them past a naive check, +# and quotes are because an unbalanced one changes where the next line starts. +FORBIDDEN_CHARS = frozenset(";|&$`><\n\r\\\"'()\x00") + +# Telnet control bytes. We only ever answer, never ask. +_IAC = 255 +_DONT, _DO, _WONT, _WILL = 254, 253, 252, 251 +_SB, _SE = 250, 240 + + +class ShellError(RuntimeError): + """The request was refused, or the drone's shell could not be reached.""" + + +class ShellUnavailable(ShellError): + """Port 23 is closed, which is the drone's normal state.""" + + +class NotAllowed(ShellError): + """The command name or an argument failed the allow-list.""" + + +@dataclass(frozen=True) +class Result: + command: str + stdout: str + + +def allowed_commands() -> list[str]: + return sorted(ALLOWED) + + +def validate(command: str, args: list[str] | tuple[str, ...] = ()) -> str: + """Check a command against the allow-list and return the line to send. + + Separate from `run` so the argv construction can be tested, and so a tool + can validate without opening a socket. + """ + if command not in ALLOWED: + raise NotAllowed( + f"'{command}' is not one of the commands this tool can run. " + f"Available: {', '.join(allowed_commands())}. The drone's shell is root, so only " + "an explicit allow-list of read-only commands is exposed." + ) + for arg in args: + if not arg: + raise NotAllowed("Empty arguments are not allowed.") + bad = sorted(set(arg) & FORBIDDEN_CHARS) + if bad: + shown = "".join(repr(c)[1:-1] for c in bad) + raise NotAllowed( + f"Argument {arg!r} contains shell metacharacters ({shown}). " + "Arguments must be literal: no pipes, redirection, substitution or quoting." + ) + return " ".join([command, *args]) + + +def run( + command: str, + args: list[str] | tuple[str, ...] = (), + *, + host: str, + port: int = TELNET_PORT, + timeout: float = DEFAULT_TIMEOUT, + max_bytes: int = DEFAULT_MAX_BYTES, +) -> Result: + """Run one allow-listed command and return its output as text. + + stderr is folded in, because when a path is wrong the message saying so is + the entire value of the call. + """ + line = validate(command, args) + + # Bracketing the command between two echoed markers is what makes the + # output findable. The alternative, trimming the shell's echo of our own + # input, needs a prompt pattern, and telnetd's pty prefixes the prompt to + # the echoed line (`/ # cat /proc/version`) so there is nothing reliable to + # match. Markers need no such guess. The nonce is what stops a command that + # prints the marker's text from ending the read early. + nonce = uuid.uuid4().hex[:12] + begin, end = f"__MCBEBOP_{nonce}_B__", f"__MCBEBOP_{nonce}_E__" + script = f"echo {begin}\n{line} 2>&1\necho {end}\n" + + raw = _converse(host, port, script.encode(), end.encode(), timeout, max_bytes) + text = raw.decode("utf-8", errors="replace") + logger.debug("ran %r on %s, %d bytes back", line, host, len(raw)) + return Result(command=line, stdout=_between(text, begin, end)) + + +def _converse(host: str, port: int, script: bytes, marker: bytes, timeout: float, max_bytes: int) -> bytes: + """Open the socket, send the script, read until the marker or the cap.""" + try: + sock = socket.create_connection((host, port), timeout=timeout) + except OSError as exc: + raise ShellUnavailable( + f"No telnet on {host}:{port} ({exc}). The Bebop 2 keeps it closed until debug mode " + "is enabled: press the power button four times, briefly. It does not survive a reboot." + ) from exc + + buf = bytearray() + try: + sock.settimeout(timeout) + sock.sendall(script) + while not _saw_marker_line(bytes(buf), marker): + try: + chunk = sock.recv(4096) + except TimeoutError as exc: + raise ShellError( + f"The drone's shell went quiet after {timeout}s without finishing. " + f"Partial output: {bytes(buf[-200:])!r}" + ) from exc + if not chunk: + break # shell exited; whatever we have is the answer + buf += _answer_negotiation(sock, chunk) + if len(buf) > max_bytes: + raise ShellError( + f"Output passed the {max_bytes}-byte cap. Narrow the command, or use " + "head/tail to take a slice of it." + ) + finally: + sock.close() + return bytes(buf) + + +def _answer_negotiation(sock: socket.socket, chunk: bytes) -> bytes: + """Strip telnet IAC sequences out of `chunk`, refusing every option. + + `DO x` gets `WONT x` and `WILL x` gets `DONT x`, which is the polite way to + say "plain text please". Subnegotiations are skipped wholesale. A sequence + A sequence split across two reads is dropped rather than reassembled, + which is safe here because busybox telnetd sends its options in one burst + before any shell output and we refuse all of them anyway. + """ + out = bytearray() + reply = bytearray() + i = 0 + while i < len(chunk): + byte = chunk[i] + if byte != _IAC: + out.append(byte) + i += 1 + continue + if i + 1 >= len(chunk): + break # truncated IAC, drop it rather than emit a stray 0xff + verb = chunk[i + 1] + if verb in (_DO, _DONT, _WILL, _WONT): + if i + 2 >= len(chunk): + break + option = chunk[i + 2] + if verb == _DO: + reply += bytes([_IAC, _WONT, option]) + elif verb == _WILL: + reply += bytes([_IAC, _DONT, option]) + i += 3 + elif verb == _SB: + end = chunk.find(bytes([_IAC, _SE]), i) + i = len(chunk) if end == -1 else end + 2 + elif verb == _IAC: + out.append(_IAC) # escaped literal 0xff + i += 2 + else: + i += 2 # a two-byte command we do not care about + if reply: + sock.sendall(bytes(reply)) + return bytes(out) + + +def _is_marker_line(stripped: str, marker: str) -> bool: + """Is this the marker being *printed*, rather than our input being echoed? + + telnetd gives the shell a pty, whose line discipline echoes everything it + reads, so each marker appears twice: once as `echo ` when the input + is read, and once alone when the command runs. Only the second is a + position in the output. A prompt may be glued to the front of either, so + the test is on the end of the line, and `echo` is what tells them apart. + """ + return stripped.endswith(marker) and "echo" not in stripped + + +def _saw_marker_line(buf: bytes, marker: bytes) -> bool: + """Has a complete marker line arrived? Used as the read sentinel. + + A partial last line does not count: the marker could still be mid-arrival, + and stopping on it would truncate the final chunk of real output. + """ + if b"\n" not in buf: + return False + text = buf.decode("utf-8", errors="replace") + complete = text.rsplit("\n", 1)[0] + return any(_is_marker_line(line.strip(), marker.decode()) for line in complete.splitlines()) + + +def _between(text: str, begin: str, end: str) -> str: + """Keep what the two markers bracket. + + The begin marker is taken as the *last* printed one before the end, and the + end as the *first* printed one, so a pty that echoes all three input lines + up front cannot make the window close before the output arrives. + """ + lines = text.splitlines() + stop = next((i for i, ln in enumerate(lines) if _is_marker_line(ln.strip(), end)), len(lines)) + start = -1 + for i in range(stop): + if _is_marker_line(lines[i].strip(), begin): + start = i + return "\n".join(lines[start + 1 : stop]).strip("\n") diff --git a/src/mcbebop/media/video.py b/src/mcbebop/media/video.py new file mode 100644 index 0000000..0ee1773 --- /dev/null +++ b/src/mcbebop/media/video.py @@ -0,0 +1,271 @@ +"""Live video from the Bebop 2, over ARStream2. + +Firmware 4.7.1 serves no RTSP at all: `192.168.42.1:554` refuses the +connection and the string `rtsp` appears nowhere in the rootfs. pyparrot's +`DroneVision` opens `rtsp://192.168.42.1/live`, which is older-firmware only, +so there is deliberately no RTSP path here. + +What it does instead is ARStream2: once `ardrone3.MediaStreaming.VideoEnable` +is sent, the drone pushes plain RTP/H.264 from `192.168.42.1:5004` to the +`arstream2_client_stream_port` the controller named in its handshake (55004), +with RTCP on 5005 to 55005. There is no URL to open. We describe our *own* +port in an SDP file and let ffmpeg bind it, which is why the ffmpeg process +has to be running before the stream is enabled: RTP is connectionless, so +packets that arrive before anything is listening are simply gone. + +The stream only flows while a controller is connected, so the caller's ARSDK +session has to stay up for the whole session. This module does not own that +session. It asks for a `Sender`, an async callable taking a command name and +an argument dict, so `media/` stays independent of `arsdk/` and `protocol/`. +The tools layer, which holds both the long-lived session and the command +index, wires them together with something like: + + async def send(name, args): + return await session.send(index.get(name), args) + +Measured on the live aircraft: 856x480 at about 30 fps, 2.5 Mbit/s. The drone +records 1080p internally but streams the 480p leg, which is the +`rec1080_stream480` setting. +""" + +from __future__ import annotations + +import asyncio +import io +import logging +import shutil +import subprocess +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + +STREAM_PORT = 55004 +PAYLOAD_TYPE = 96 # dynamic; the drone's SSRC announces H.264 at 90 kHz +VIDEO_ENABLE = "ardrone3.MediaStreaming.VideoEnable" +EXPOSURE = "ardrone3.PictureSettings.ExpositionSelection" + +# The aircraft accepts -1.5..1.5 here; outside that it ignores the command. +EXPOSURE_RANGE = (-1.5, 1.5) +# What the drone reports its current exposure as, in the telemetry dict. +EXPOSURE_STATE_KEY = "ExpositionChanged_value" + +# ffmpeg will not touch rtp/udp from an SDP unless they are whitelisted. +_PROTOCOLS = ["-protocol_whitelist", "file,rtp,udp"] +_QUIET = ["-hide_banner", "-loglevel", "warning"] +# Without these, ffmpeg buffers for smoothness and the feed runs seconds late. +_LOW_LATENCY = ["-fflags", "nobuffer", "-flags", "low_delay"] + +#: An async callable `(command_name, args) -> result`. See the module docstring. +Sender = Callable[[str, dict[str, Any]], Awaitable[Any]] +#: A synchronous telemetry reader, `(keys) -> values`, matching `Session.state`. +StateReader = Callable[[list[str]], dict[str, Any]] + + +class VideoUnavailable(RuntimeError): + """ffmpeg/ffplay is missing, or the drone refused to start streaming.""" + + +def sdp_text(port: int = STREAM_PORT, payload_type: int = PAYLOAD_TYPE) -> str: + """The SDP that describes our own receiving port. + + `c=IN IP4 0.0.0.0` is what makes ffmpeg bind rather than connect; naming + the drone's address here would make it wait for a stream it has to send. + """ + return ( + "v=0\n" + "o=- 0 0 IN IP4 127.0.0.1\n" + "s=Parrot Bebop 2\n" + "c=IN IP4 0.0.0.0\n" + "t=0 0\n" + f"m=video {port} RTP/AVP {payload_type}\n" + f"a=rtpmap:{payload_type} H264/90000\n" + ) + + +def write_sdp(path: Path, port: int = STREAM_PORT) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(sdp_text(port), encoding="utf-8") + return path + + +@dataclass(frozen=True) +class Sink: + """How to consume the stream: a command plus what it produces.""" + + argv: list[str] + tool: str + describe: str + output: Path | None = None + + +def view(sdp: Path, title: str = "Bebop 2") -> Sink: + return Sink( + ["ffplay", *_QUIET, *_PROTOCOLS, *_LOW_LATENCY, "-framedrop", "-window_title", title, "-i", str(sdp)], + "ffplay", + "live window", + ) + + +def record(sdp: Path, out: Path, seconds: float | None = None) -> Sink: + # -c copy keeps the drone's own H.264 untouched: no re-encode, no quality + # loss, and little CPU. The container just needs a timebase. + argv = [ + "ffmpeg", *_QUIET, "-y", *_PROTOCOLS, + "-use_wallclock_as_timestamps", "1", "-i", str(sdp), "-c", "copy", + ] # fmt: skip + if seconds: + argv += ["-t", f"{seconds:g}"] + argv += ["-movflags", "+faststart", str(out)] + return Sink(argv, "ffmpeg", f"recording to {out}", out) + + +def snapshot(sdp: Path, out: Path, settle_seconds: float = 4.0) -> Sink: + # Two reasons to discard the start: the first frames reference an SPS/PPS + # we have not seen yet ("non-existing PPS 0 referenced", which recovers on + # its own and is not an error), and the camera's auto exposure needs a few + # seconds or the frame comes out nearly black. + # -update is required for a single image (image2 otherwise wants %03d). + argv = [ + "ffmpeg", *_QUIET, "-y", *_PROTOCOLS, "-i", str(sdp), + "-ss", f"{settle_seconds:g}", "-frames:v", "1", "-update", "1", str(out), + ] # fmt: skip + return Sink(argv, "ffmpeg", f"single frame to {out}", out) + + +async def enable(send: Sender) -> None: + """Ask the drone to start pushing RTP. Bind the receiver first.""" + await send(VIDEO_ENABLE, {"enable": 1}) + + +async def disable(send: Sender) -> None: + """Stop the stream, so the aircraft is left as we found it.""" + await send(VIDEO_ENABLE, {"enable": 0}) + + +async def set_exposure(send: Sender, value: float) -> None: + low, high = EXPOSURE_RANGE + if not low <= value <= high: + raise ValueError(f"exposure must be between {low} and {high}, got {value}") + await send(EXPOSURE, {"value": float(value)}) + + +def downscale(data: bytes, max_width: int) -> bytes: + """Shrink an encoded frame to `max_width`, keeping its format. + + A full 856x480 frame is a few hundred kilobytes of base64 by the time it + reaches a model, which is most of a context window spent on pixels nobody + asked for. Already-narrow images come back untouched, bytes for bytes, so + re-encoding never degrades something that did not need resizing. + """ + if max_width <= 0: + raise ValueError(f"max_width must be positive, got {max_width}") + + from PIL import Image + + with Image.open(io.BytesIO(data)) as im: + if im.width <= max_width: + return data + fmt = im.format or "PNG" + height = max(1, round(im.height * max_width / im.width)) + shrunk = im.resize((max_width, height), Image.Resampling.LANCZOS) + buf = io.BytesIO() + # JPEG cannot hold an alpha channel, which a PNG snapshot may carry. + if fmt == "JPEG" and shrunk.mode not in ("RGB", "L"): + shrunk = shrunk.convert("RGB") + shrunk.save(buf, format=fmt) + logger.debug("downscaled %s from %d to %d px wide", fmt, im.width, max_width) + return buf.getvalue() + + +class StreamSession: + """Run a sink while the drone streams, and always turn the stream off. + + An async context manager, because enabling the stream goes through the + caller's ARSDK session. The ordering inside `__aenter__` is the part that + matters and the part that is easy to get wrong: the subprocess binds the + port *before* VideoEnable is sent. + """ + + def __init__( + self, + sink: Sink, + send: Sender, + *, + exposure: float | None = None, + read_state: StateReader | None = None, + ) -> None: + self.sink = sink + self.send = send + self.exposure = exposure + self.read_state = read_state + self.prior_exposure: float | None = None + self.proc: subprocess.Popen | None = None + + async def __aenter__(self) -> StreamSession: + if shutil.which(self.sink.tool) is None: + raise VideoUnavailable(f"{self.sink.tool} not found; install ffmpeg") + # Bind the port first: RTP is connectionless, so anything that arrives + # before the sink is listening is simply lost. + self.proc = subprocess.Popen(self.sink.argv) + logger.debug("started %s: %s", self.sink.tool, self.sink.describe) + exposure_set = False + try: + if self.exposure is not None: + if self.read_state is not None: + # remember what it was so the drone is left as we found it + prior = self.read_state([EXPOSURE_STATE_KEY]).get(EXPOSURE_STATE_KEY) + self.prior_exposure = float(prior) if prior is not None else 0.0 + await set_exposure(self.send, self.exposure) + exposure_set = True + await enable(self.send) + except BaseException: + # `__aexit__` never runs for a failed `__aenter__`, so an exposure + # we already changed has to be put back here or the next snapshot + # inherits it. + if exposure_set and self.prior_exposure is not None: + try: + await set_exposure(self.send, self.prior_exposure) + except Exception: + logger.warning("could not restore exposure after a failed start", exc_info=True) + await self._stop_proc() + raise + return self + + async def __aexit__(self, *exc: object) -> None: + try: + await disable(self.send) + if self.prior_exposure is not None: + await set_exposure(self.send, self.prior_exposure) + except Exception: + # Teardown must not mask whatever the body raised, and a stream + # left enabled is harmless once the controller link drops. + logger.warning("could not restore the drone's video state", exc_info=True) + finally: + await self._stop_proc() + + async def _stop_proc(self) -> None: + proc, self.proc = self.proc, None + if proc is None or proc.poll() is not None: + return + proc.terminate() + try: + await asyncio.to_thread(proc.wait, 5) + except subprocess.TimeoutExpired: + logger.warning("%s ignored SIGTERM; killing it", self.sink.tool) + proc.kill() + + async def wait(self, seconds: float | None = None) -> int | None: + """Block until the sink exits or the timeout elapses. + + Returns the exit status, or None if it is still running. The wait runs + off the event loop so the rest of the server keeps answering. + """ + if self.proc is None: + return None + try: + return await asyncio.to_thread(self.proc.wait, seconds) + except subprocess.TimeoutExpired: + return None diff --git a/tests/test_ftp.py b/tests/test_ftp.py new file mode 100644 index 0000000..ecd9603 --- /dev/null +++ b/tests/test_ftp.py @@ -0,0 +1,242 @@ +"""FTP: the area-to-port map, path containment, and the size cap. + +ftplib is replaced with a fake; nothing opens a socket. +""" + +import ftplib +from typing import ClassVar + +import pytest + +from mcbebop.files import ftp + +# --- the area map ------------------------------------------------------------ + + +def test_the_three_read_only_areas_map_to_the_ports_found_on_the_aircraft(): + assert {a.name: a.port for a in ftp.areas()} == {"media": 21, "flightplans": 61, "logs": 21} + + +def test_the_firmware_write_port_is_not_reachable_through_the_api(): + # Port 51 serves /update as root. It is how firmware is pushed, it has no + # read use case, and the drone's Wi-Fi is open. Exposing it would hand the + # aircraft's weakest point to anything that can call a tool. + assert ftp.FIRMWARE_WRITE_PORT == 51 + assert 51 not in {a.port for a in ftp.areas()} + for name in ("update", "firmware", "51"): + with pytest.raises(ftp.FtpError): + ftp.resolve_area(name) + + +def test_an_unknown_area_says_what_is_available(): + with pytest.raises(ftp.FtpError, match="flightplans"): + ftp.resolve_area("nope") + + +def test_logs_is_the_media_ftpd_narrowed_to_the_blackbox_tree(): + logs = ftp.resolve_area("logs") + media = ftp.resolve_area("media") + assert logs.port == media.port == 21 + assert logs.prefix == "internal_000/Debug/current" + assert media.prefix == "" + + +def test_every_area_describes_itself(): + assert all(a.describe for a in ftp.areas()) + + +# --- path containment -------------------------------------------------------- + + +def test_a_logs_path_stays_inside_the_debug_tree(): + area = ftp.resolve_area("logs") + assert ftp._safe_path(area, "") == "internal_000/Debug/current" + assert ftp._safe_path(area, "boot.log") == "internal_000/Debug/current/boot.log" + + +def test_dot_dot_cannot_climb_out_of_an_area(): + area = ftp.resolve_area("logs") + for path in ("../../..", "a/../../b", ".."): + with pytest.raises(ftp.FtpError, match=r"\.\."): + ftp._safe_path(area, path) + + +def test_an_absolute_path_is_refused_rather_than_silently_rebased(): + with pytest.raises(ftp.FtpError, match="relative"): + ftp._safe_path(ftp.resolve_area("media"), "/data/ftp/internal_000") + + +def test_an_empty_media_path_lists_the_area_root(): + assert ftp._safe_path(ftp.resolve_area("media"), "") == "." + + +def test_redundant_separators_and_dots_collapse(): + area = ftp.resolve_area("media") + assert ftp._safe_path(area, "internal_000//./media/") == "internal_000/media" + + +# --- listing ----------------------------------------------------------------- + + +UNIX_LINE = "-rw-r--r-- 1 root root 1048576 Jan 1 00:00 flightPlan.mavlink" +DIR_LINE = "drwxr-xr-x 2 root root 4096 Jan 1 00:00 internal_000" + + +def test_a_unix_list_line_parses(): + entry = ftp._parse_line(UNIX_LINE) + assert entry.name == "flightPlan.mavlink" + assert entry.size == 1048576 + assert entry.is_dir is False + + +def test_a_directory_line_is_marked_as_one(): + assert ftp._parse_line(DIR_LINE).is_dir is True + + +def test_an_unrecognised_line_still_yields_a_name(): + entry = ftp._parse_line("weird-output.txt") + assert entry.name == "weird-output.txt" + assert entry.size is None + assert entry.raw == "weird-output.txt" + + +# --- a fake ftpd ------------------------------------------------------------- + + +class FakeFTP: + """Just enough ftplib surface, recording what was asked of it. + + `FakeFTP.config` is what a test uses to describe the ftpd it wants, since + ftplib.FTP is constructed inside the module under test. + """ + + instances: ClassVar[list["FakeFTP"]] = [] + config: ClassVar[dict] = {} + + def __init__(self, timeout=None): + self.timeout = timeout + self.connected_to = None + self.logged_in = False + self.commands: list[str] = [] + self.quit_called = False + self.listing = FakeFTP.config.get("listing", [DIR_LINE, UNIX_LINE]) + self.content = FakeFTP.config.get("content", b"x" * 64) + self.declared_size = FakeFTP.config.get("declared_size") + self.size_raises = FakeFTP.config.get("size_raises", False) + FakeFTP.instances.append(self) + + def connect(self, host, port): + self.connected_to = (host, port) + + def login(self, *a): + self.logged_in = True + + def retrlines(self, cmd, callback): + self.commands.append(cmd) + for line in self.listing: + callback(line) + + def retrbinary(self, cmd, callback): + self.commands.append(cmd) + for i in range(0, len(self.content), 16): + callback(self.content[i : i + 16]) + + def voidcmd(self, cmd): + self.commands.append(cmd) + + def size(self, path): + if self.size_raises: + raise ftplib.error_perm("550 SIZE not understood") + return self.declared_size + + def quit(self): + self.quit_called = True + + def close(self): + pass + + +@pytest.fixture +def fake_ftp(monkeypatch): + FakeFTP.instances = [] + FakeFTP.config = {} + monkeypatch.setattr(ftp.ftplib, "FTP", FakeFTP) + return FakeFTP + + +def test_list_dir_connects_to_the_area_port_and_logs_in_anonymously(fake_ftp): + entries = ftp.list_dir("flightplans", host="192.168.42.1") + conn = fake_ftp.instances[0] + assert conn.connected_to == ("192.168.42.1", 61) + assert conn.logged_in + assert conn.timeout == ftp.DEFAULT_TIMEOUT + assert [e.name for e in entries] == ["internal_000", "flightPlan.mavlink"] + assert conn.quit_called + + +def test_list_dir_scopes_the_logs_area_in_the_command_it_sends(fake_ftp): + ftp.list_dir("logs", host="h") + assert fake_ftp.instances[0].commands[0] == "LIST internal_000/Debug/current" + + +def test_a_refused_connection_explains_the_access_point(fake_ftp, monkeypatch): + def refuse(self, host, port): + raise OSError("connection refused") + + monkeypatch.setattr(FakeFTP, "connect", refuse) + with pytest.raises(ftp.FtpError, match="access point"): + ftp.list_dir("media", host="h") + + +# --- fetching ---------------------------------------------------------------- + + +def test_fetch_streams_to_capture_dir_and_returns_the_path(fake_ftp, tmp_path): + dest = ftp.fetch("flightplans", "flightPlan.mavlink", host="h", capture_dir=tmp_path) + assert dest == tmp_path / "flightplans" / "flightPlan.mavlink" + assert dest.read_bytes() == b"x" * 64 + assert "RETR flightPlan.mavlink" in fake_ftp.instances[0].commands + + +def test_fetch_asks_for_the_size_in_binary_mode_first(fake_ftp, tmp_path): + # busybox will not answer SIZE in ASCII mode, so TYPE I has to come first. + ftp.fetch("media", "internal_000/a.jpg", host="h", capture_dir=tmp_path) + cmds = fake_ftp.instances[0].commands + assert cmds[0] == "TYPE I" + + +def test_a_file_the_drone_declares_as_oversized_is_never_transferred(fake_ftp, tmp_path): + fake_ftp.config["declared_size"] = 50_000_000 + with pytest.raises(ftp.TooLarge, match="cap"): + ftp.fetch("media", "internal_000/big.mp4", host="h", capture_dir=tmp_path) + assert not any(c.startswith("RETR") for c in fake_ftp.instances[0].commands) + assert not (tmp_path / "media" / "big.mp4").exists() + + +def test_a_cap_passed_mid_transfer_aborts_and_removes_the_partial_file(fake_ftp, tmp_path): + # The second line of defence: an ftpd that will not answer SIZE still + # cannot stream a video into memory, because the callback counts bytes. + fake_ftp.config.update(content=b"y" * 4096, size_raises=True) + with pytest.raises(ftp.TooLarge, match="mid-transfer"): + ftp.fetch("media", "internal_000/big.mp4", host="h", capture_dir=tmp_path, max_bytes=100) + assert not (tmp_path / "media" / "big.mp4").exists() + + +def test_fetch_refuses_a_directory(fake_ftp, tmp_path): + with pytest.raises(ftp.FtpError, match="file path"): + ftp.fetch("media", "", host="h", capture_dir=tmp_path) + + +def test_fetch_cannot_be_talked_out_of_its_area(fake_ftp, tmp_path): + with pytest.raises(ftp.FtpError): + ftp.fetch("logs", "../../internal_000/video.mp4", host="h", capture_dir=tmp_path) + + +def test_a_server_error_during_fetch_leaves_no_stub_file(fake_ftp, tmp_path, monkeypatch): + def fail(self, cmd, callback): + raise ftplib.error_perm("550 No such file") + + monkeypatch.setattr(FakeFTP, "retrbinary", fail) + with pytest.raises(ftp.FtpError, match="Could not fetch"): + ftp.fetch("media", "internal_000/ghost.jpg", host="h", capture_dir=tmp_path) + assert not (tmp_path / "media" / "ghost.jpg").exists() diff --git a/tests/test_shell.py b/tests/test_shell.py new file mode 100644 index 0000000..baf1cd5 --- /dev/null +++ b/tests/test_shell.py @@ -0,0 +1,378 @@ +"""Shell: the allow-list, metacharacter refusal, and the telnet handshake. + +The socket is a fake. The point of most of these tests is that nothing +reaches it: validation happens before a byte is sent. +""" + +import socket + +import pytest + +from mcbebop.files import shell + +# --- the allow-list ---------------------------------------------------------- + + +def test_only_read_only_commands_are_allowed(): + assert shell.allowed_commands() == [ + "bcmwl", + "cat", + "df", + "dmesg", + "getprop", + "head", + "ls", + "mount", + "ps", + "tail", + "uptime", + ] + + +@pytest.mark.parametrize( + "name", + [ + "rm", + "sh", + "bash", + "dd", + "mv", + "cp", + "chmod", + "reboot", + "telnetd", + "echo", + "eval", + "python", + "CAT", + "ls ", + ], +) +def test_a_command_outside_the_allow_list_is_refused(name): + # Deny-lists on a root shell leak; there is always another spelling. So + # the only names that work are the ones explicitly listed. + with pytest.raises(shell.NotAllowed, match="not one of"): + shell.validate(name) + + +def test_the_refusal_lists_what_is_available(): + with pytest.raises(shell.NotAllowed, match="getprop"): + shell.validate("rm") + + +@pytest.mark.parametrize("meta", [";", "|", "&", "$", "`", ">", "<", "\n", "\r", "\\", '"', "'", "(", ")"]) +def test_an_argument_carrying_a_shell_metacharacter_is_refused(meta): + with pytest.raises(shell.NotAllowed, match="metacharacter"): + shell.validate("cat", [f"/proc/version{meta}"]) + + +@pytest.mark.parametrize( + "arg", + [ + "/etc/passwd; rm -rf /", + "$(reboot)", + "`id`", + "/dev/null > /bin/sh", + "a | sh", + "x && reboot", + "a\necho hi", + "\\;", + ], +) +def test_the_classic_injections_never_get_past_validate(arg): + with pytest.raises(shell.NotAllowed): + shell.validate("cat", [arg]) + + +def test_an_empty_argument_is_refused(): + with pytest.raises(shell.NotAllowed, match="Empty"): + shell.validate("ls", [""]) + + +def test_flags_and_ordinary_paths_are_fine(): + assert shell.validate("ls", ["-l", "/data/ftp"]) == "ls -l /data/ftp" + assert shell.validate("getprop", ["ro.parrot.build.version"]) == "getprop ro.parrot.build.version" + assert shell.validate("tail", ["-n", "40", "/var/log/messages"]) == "tail -n 40 /var/log/messages" + assert shell.validate("uptime") == "uptime" + + +def test_validation_happens_before_the_socket_opens(monkeypatch): + def explode(*a, **kw): + raise AssertionError("must not connect for a refused command") + + monkeypatch.setattr(shell.socket, "create_connection", explode) + with pytest.raises(shell.NotAllowed): + shell.run("rm", ["-rf", "/"], host="192.168.42.1") + with pytest.raises(shell.NotAllowed): + shell.run("cat", ["/etc/passwd; reboot"], host="192.168.42.1") + + +# --- a fake telnetd ---------------------------------------------------------- + + +class FakeSocket: + """Replays a scripted server side and records what we sent.""" + + def __init__(self, chunks): + self.chunks = list(chunks) + self.sent = bytearray() + self.closed = False + self.timeout = None + + def settimeout(self, t): + self.timeout = t + + def sendall(self, data): + self.sent += data + + def recv(self, n): + if not self.chunks: + return b"" + chunk = self.chunks.pop(0) + if chunk is TimeoutError: + raise TimeoutError("timed out") + return chunk + + def close(self): + self.closed = True + + +@pytest.fixture +def fake_telnet(): + """Somewhere for a scripted connect() to record what it saw.""" + return {} + + +def _markers(sock: FakeSocket) -> tuple[str, str]: + """Dig the two markers out of what the module sent, so a fake can echo them.""" + lines = sock.sent.decode().strip().splitlines() + return lines[0].split()[-1], lines[-1].split()[-1] + + +def _pty(holder, *, output=b"", banner=b"", trailing=b"", negotiate=b""): + """A telnetd that behaves like a real one: pty echo first, output after. + + The echo of all three input lines arrives in one burst *before* anything + runs, so a reader that stopped at the first sight of the end marker would + return nothing at all. That is the case this fake exists to reproduce. + """ + + def connect(address, timeout=None): + holder["address"] = address + holder["timeout"] = timeout + + class Responder(FakeSocket): + def recv(self, n): + if not self.sent or self.chunks == ["done"]: + return b"" + begin, end = _markers(self) + echo = self.sent.decode() + self.chunks = ["done"] + return ( + negotiate + + banner + # the pty echoing our input, prompt glued to the front + + b"".join(f"/ # {ln}\n".encode() for ln in echo.splitlines()) + + f"{begin}\n".encode() + + (output + b"\n" if output else b"") + + f"/ # {end}\n".encode() + + trailing + ) + + sock = Responder([]) + holder["sock"] = sock + return sock + + return connect + + +def test_a_command_runs_and_its_output_comes_back(monkeypatch, fake_telnet): + monkeypatch.setattr(shell.socket, "create_connection", _pty(fake_telnet, output=b"Linux version 3.4.11")) + result = shell.run("cat", ["/proc/version"], host="192.168.42.1") + assert result.command == "cat /proc/version" + assert result.stdout == "Linux version 3.4.11" + assert fake_telnet["address"] == ("192.168.42.1", 23) + assert fake_telnet["sock"].closed + + +def test_the_pty_echo_of_our_own_input_never_reaches_the_caller(monkeypatch, fake_telnet): + monkeypatch.setattr(shell.socket, "create_connection", _pty(fake_telnet, output=b"3.4.11")) + out = shell.run("cat", ["/proc/version"], host="h").stdout + assert out == "3.4.11" + assert "cat /proc/version" not in out + assert "MCBEBOP" not in out + + +def test_a_login_banner_before_the_markers_is_dropped(monkeypatch, fake_telnet): + monkeypatch.setattr( + shell.socket, + "create_connection", + _pty(fake_telnet, banner=b"BusyBox v1.20.2 built-in shell\n\n", output=b"4.7.1"), + ) + assert shell.run("getprop", ["ro.parrot.build.version"], host="h").stdout == "4.7.1" + + +def test_anything_after_the_end_marker_is_discarded(monkeypatch, fake_telnet): + monkeypatch.setattr( + shell.socket, "create_connection", _pty(fake_telnet, output=b"real output", trailing=b"/ # \nnoise\n") + ) + assert shell.run("uptime", host="h").stdout == "real output" + + +def test_multi_line_output_keeps_its_blank_lines_and_indentation(monkeypatch, fake_telnet): + monkeypatch.setattr(shell.socket, "create_connection", _pty(fake_telnet, output=b"one\n\n three")) + assert shell.run("dmesg", host="h").stdout == "one\n\n three" + + +def test_a_command_with_no_output_returns_an_empty_string(monkeypatch, fake_telnet): + monkeypatch.setattr(shell.socket, "create_connection", _pty(fake_telnet)) + assert shell.run("ls", ["/nonexistent-but-quiet"], host="h").stdout == "" + + +def test_the_command_we_send_is_one_simple_command(monkeypatch, fake_telnet): + monkeypatch.setattr(shell.socket, "create_connection", _pty(fake_telnet, output=b"ok")) + shell.run("ls", ["-l", "/data/ftp"], host="h") + lines = fake_telnet["sock"].sent.decode().strip().splitlines() + # Three lines: open marker, the command, close marker. stderr is merged so + # a failure explains itself; nothing else is added. + assert len(lines) == 3 + assert lines[1] == "ls -l /data/ftp 2>&1" + assert lines[0].startswith("echo __MCBEBOP_") and lines[2].startswith("echo __MCBEBOP_") + + +def test_a_closed_port_says_how_to_open_it(monkeypatch): + def refuse(address, timeout=None): + raise ConnectionRefusedError("connection refused") + + monkeypatch.setattr(shell.socket, "create_connection", refuse) + with pytest.raises(shell.ShellUnavailable, match="four times"): + shell.run("uptime", host="192.168.42.1") + + +def test_a_silent_shell_times_out_rather_than_hanging(monkeypatch, fake_telnet): + def connect(address, timeout=None): + sock = FakeSocket([TimeoutError]) + fake_telnet["sock"] = sock + return sock + + monkeypatch.setattr(shell.socket, "create_connection", connect) + with pytest.raises(shell.ShellError, match="went quiet"): + shell.run("dmesg", host="h", timeout=0.01) + + +def test_a_shell_that_exits_early_returns_what_arrived(monkeypatch, fake_telnet): + def connect(address, timeout=None): + sock = FakeSocket([b"partial output\n", b""]) + fake_telnet["sock"] = sock + return sock + + monkeypatch.setattr(shell.socket, "create_connection", connect) + assert shell.run("dmesg", host="h").stdout == "partial output" + + +def test_a_flood_of_output_is_capped(monkeypatch, fake_telnet): + def connect(address, timeout=None): + sock = FakeSocket([b"z" * 4096] * 100) + fake_telnet["sock"] = sock + return sock + + monkeypatch.setattr(shell.socket, "create_connection", connect) + with pytest.raises(shell.ShellError, match="cap"): + shell.run("cat", ["/dev/urandom"], host="h", max_bytes=1024) + + +# --- telnet negotiation ------------------------------------------------------ + + +IAC, DO, DONT, WILL, WONT, SB, SE = 255, 253, 254, 251, 252, 250, 240 +ECHO_OPT, SGA = 1, 3 + + +def test_every_option_the_server_offers_is_declined(): + sock = FakeSocket([]) + data = shell._answer_negotiation(sock, bytes([IAC, DO, ECHO_OPT, IAC, WILL, SGA]) + b"hello") + assert data == b"hello" # control bytes never reach the caller + assert bytes(sock.sent) == bytes([IAC, WONT, ECHO_OPT, IAC, DONT, SGA]) + + +def test_a_subnegotiation_block_is_skipped_whole(): + sock = FakeSocket([]) + payload = bytes([IAC, SB, 24, 0, 65, 66, IAC, SE]) + b"after" + assert shell._answer_negotiation(sock, payload) == b"after" + assert bytes(sock.sent) == b"" + + +def test_an_escaped_literal_ff_survives(): + sock = FakeSocket([]) + assert shell._answer_negotiation(sock, b"a" + bytes([IAC, IAC]) + b"b") == b"a\xffb" + + +def test_a_truncated_sequence_does_not_leak_a_stray_control_byte(): + sock = FakeSocket([]) + assert shell._answer_negotiation(sock, b"ok" + bytes([IAC, DO])) == b"ok" + assert shell._answer_negotiation(sock, b"ok" + bytes([IAC])) == b"ok" + + +def test_negotiation_is_stripped_out_of_real_output(monkeypatch, fake_telnet): + monkeypatch.setattr( + shell.socket, + "create_connection", + _pty(fake_telnet, negotiate=bytes([IAC, DO, ECHO_OPT]), output=b"3.4.11"), + ) + assert shell.run("cat", ["/proc/version"], host="h").stdout == "3.4.11" + + +def test_the_module_never_imports_telnetlib(): + # Removed in Python 3.13, which is why the protocol lives here. + with pytest.raises(ImportError): + __import__("telnetlib") + assert "telnetlib" not in dir(shell) + + +def test_we_use_the_stdlib_socket_module_directly(): + assert shell.socket is socket + + +# --- bracketing a realistic busybox transcript ------------------------------- + + +def test_a_busybox_transcript_is_reduced_to_the_output(): + begin, end = "__MCBEBOP_dead_B__", "__MCBEBOP_dead_E__" + transcript = ( + "BusyBox v1.20.2 built-in shell\n" + "\n" + f"/ # echo {begin}\n" + "/ # getprop ro.parrot.build.version 2>&1\n" + f"/ # echo {end}\n" + f"{begin}\n" + "4.7.1\n" + f"/ # {end}\n" + "/ # " + ) + assert shell._between(transcript, begin, end) == "4.7.1" + + +def test_output_that_happens_to_look_like_a_prompt_is_kept(): + begin, end = "__B_B__", "__B_E__" + assert shell._between(f"{begin}\ndata\n#\nmore\n{end}\n", begin, end) == "data\n#\nmore" + + +def test_without_a_begin_marker_everything_up_to_the_end_is_kept(): + # An echo-less shell that somehow skipped the opening marker should still + # yield output rather than nothing. + assert shell._between("4.7.1\n__E__\n", "__B__", "__E__") == "4.7.1" + + +def test_a_marker_line_is_told_apart_from_its_own_echo(): + assert shell._is_marker_line("/ # __M__", "__M__") + assert shell._is_marker_line("__M__", "__M__") + assert not shell._is_marker_line("echo __M__", "__M__") + assert not shell._is_marker_line("/ # echo __M__", "__M__") + + +def test_the_read_sentinel_waits_for_a_complete_line(): + # A marker still arriving must not stop the read, or the last chunk of + # real output is lost. + assert not shell._saw_marker_line(b"out\n__M_", b"__M__") + assert not shell._saw_marker_line(b"echo __M__\n", b"__M__") + assert shell._saw_marker_line(b"out\n__M__\n", b"__M__") diff --git a/tests/test_video.py b/tests/test_video.py new file mode 100644 index 0000000..f7d2cf2 --- /dev/null +++ b/tests/test_video.py @@ -0,0 +1,357 @@ +"""Video: SDP text, ffmpeg argv, the downscale helper, and the start ordering. + +Nothing here spawns ffmpeg or touches a drone. +""" + +import io +import subprocess +from pathlib import Path + +import pytest +from PIL import Image + +from mcbebop.media import video + + +def test_sdp_describes_our_own_port_not_the_drone(): + # c= must be the wildcard: naming the drone makes ffmpeg wait to be sent + # to, which is backwards. The port is ours, the one we named in the + # handshake, and ffmpeg binds it. + text = video.sdp_text() + assert "c=IN IP4 0.0.0.0" in text + assert f"m=video {video.STREAM_PORT} RTP/AVP 96" in text + assert "a=rtpmap:96 H264/90000" in text + assert text.startswith("v=0\n") + assert "rtsp" not in text.lower() + + +def test_sdp_lines_are_in_order_and_complete(): + keys = [line.split("=", 1)[0] for line in video.sdp_text().strip().splitlines()] + assert keys == ["v", "o", "s", "c", "t", "m", "a"] + + +def test_sdp_port_is_overridable(): + assert "m=video 60000 RTP/AVP 96" in video.sdp_text(port=60000) + + +def test_write_sdp_creates_parents(tmp_path): + out = video.write_sdp(tmp_path / "nested" / "bebop.sdp") + assert out.read_text() == video.sdp_text() + + +@pytest.fixture +def sdp(tmp_path): + return tmp_path / "bebop.sdp" + + +def test_every_sink_whitelists_rtp_and_udp(sdp, tmp_path): + # ffmpeg refuses to open rtp/udp referenced from a file-based SDP unless + # they are explicitly whitelisted, and the failure reads like a bad file. + sinks = [ + video.view(sdp), + video.record(sdp, tmp_path / "out.mp4"), + video.snapshot(sdp, tmp_path / "out.png"), + ] + for sink in sinks: + argv = sink.argv + assert "-protocol_whitelist" in argv + assert argv[argv.index("-protocol_whitelist") + 1] == "file,rtp,udp" + + +def test_record_copies_the_stream_rather_than_re_encoding(sdp, tmp_path): + argv = video.record(sdp, tmp_path / "out.mp4").argv + assert argv[argv.index("-c") + 1] == "copy" + assert "-t" not in argv + assert argv[-1] == str(tmp_path / "out.mp4") + + +def test_record_with_a_duration_passes_t(sdp, tmp_path): + argv = video.record(sdp, tmp_path / "out.mp4", seconds=12.5).argv + assert argv[argv.index("-t") + 1] == "12.5" + + +def test_snapshot_skips_the_settle_period_and_asks_for_one_image(sdp, tmp_path): + # -ss discards the frames that reference an SPS/PPS we have not seen yet + # ("non-existing PPS 0") and lets auto-exposure settle. -update is what + # lets a single filename work at all; image2 otherwise demands %03d. + argv = video.snapshot(sdp, tmp_path / "f.png", settle_seconds=3).argv + assert argv[argv.index("-ss") + 1] == "3" + assert argv[argv.index("-frames:v") + 1] == "1" + assert argv[argv.index("-update") + 1] == "1" + # -ss has to come after -i here: we are discarding received frames, not + # seeking in a file. + assert argv.index("-i") < argv.index("-ss") + + +def test_snapshot_settle_default_is_not_zero(sdp, tmp_path): + argv = video.snapshot(sdp, tmp_path / "f.png").argv + assert float(argv[argv.index("-ss") + 1]) > 0 + + +def test_view_uses_ffplay_and_low_latency_flags(sdp): + sink = video.view(sdp, title="hello") + assert sink.tool == "ffplay" + assert sink.argv[0] == "ffplay" + assert "-nobuffer" not in sink.argv # it is a value, not a flag + assert sink.argv[sink.argv.index("-fflags") + 1] == "nobuffer" + assert sink.argv[sink.argv.index("-window_title") + 1] == "hello" + + +def test_no_sink_opens_a_stream_url(sdp, tmp_path): + # Firmware 4.7.1 refuses port 554 entirely, so an rtsp:// URL anywhere + # here would be a regression to what pyparrot does. Match on the scheme + # rather than the bare word: pytest's own tmp_path is named after the test. + sinks = (video.view(sdp), video.record(sdp, tmp_path / "a.mp4"), video.snapshot(sdp, tmp_path / "a.png")) + for sink in sinks: + assert not any(part.startswith(("rtsp://", "rtmp://", "http://")) for part in sink.argv) + assert not any(":554" in part for part in sink.argv) + + +# --- downscale --------------------------------------------------------------- + + +def _png(width: int, height: int) -> bytes: + buf = io.BytesIO() + Image.new("RGB", (width, height), (40, 90, 140)).save(buf, format="PNG") + return buf.getvalue() + + +def _jpeg(width: int, height: int) -> bytes: + buf = io.BytesIO() + Image.new("RGB", (width, height), (40, 90, 140)).save(buf, format="JPEG") + return buf.getvalue() + + +def test_downscale_shrinks_and_keeps_aspect_ratio(): + out = video.downscale(_png(856, 480), 640) + with Image.open(io.BytesIO(out)) as im: + assert im.width == 640 + assert im.height == round(480 * 640 / 856) + assert im.format == "PNG" + + +def test_downscale_leaves_a_narrow_image_byte_for_byte(): + data = _png(320, 180) + assert video.downscale(data, 640) is data + + +def test_downscale_at_exactly_max_width_is_a_passthrough(): + data = _png(640, 360) + assert video.downscale(data, 640) is data + + +def test_downscale_preserves_jpeg_format(): + out = video.downscale(_jpeg(856, 480), 200) + with Image.open(io.BytesIO(out)) as im: + assert im.format == "JPEG" + assert im.width == 200 + + +def test_downscale_handles_rgba_into_jpeg(): + buf = io.BytesIO() + Image.new("RGBA", (856, 480), (1, 2, 3, 255)).save(buf, format="PNG") + out = video.downscale(buf.getvalue(), 100) + with Image.open(io.BytesIO(out)) as im: + assert im.width == 100 + + +def test_downscale_actually_saves_bytes(): + big = _png(856, 480) + assert len(video.downscale(big, 160)) < len(big) + + +def test_downscale_rejects_a_nonsense_width(): + with pytest.raises(ValueError): + video.downscale(_png(856, 480), 0) + + +# --- StreamSession ----------------------------------------------------------- + + +class FakeProc: + def __init__(self, argv): + self.argv = argv + self.terminated = False + self.killed = False + self._returncode = None + + def poll(self): + return self._returncode + + def terminate(self): + self.terminated = True + self._returncode = 0 + + def kill(self): + self.killed = True + self._returncode = -9 + + def wait(self, timeout=None): + return 0 + + +@pytest.fixture +def harness(monkeypatch, tmp_path): + """Record the order of (subprocess start, command send) events.""" + events: list[tuple] = [] + procs: list[FakeProc] = [] + + def fake_popen(argv, *a, **kw): + events.append(("popen", argv)) + proc = FakeProc(argv) + procs.append(proc) + return proc + + monkeypatch.setattr(video.subprocess, "Popen", fake_popen) + monkeypatch.setattr(video.shutil, "which", lambda tool: f"/usr/bin/{tool}") + + async def send(name, args): + events.append(("send", name, args)) + return {"ok": True} + + return events, procs, send + + +async def test_the_receiver_binds_before_the_stream_is_enabled(harness, tmp_path): + # RTP is connectionless. Enable the stream first and the opening packets, + # which carry the SPS/PPS, hit a closed port and are gone. + events, _, send = harness + sink = video.snapshot(tmp_path / "s.sdp", tmp_path / "f.png") + async with video.StreamSession(sink, send): + pass + assert events[0][0] == "popen" + assert events[1] == ("send", video.VIDEO_ENABLE, {"enable": 1}) + + +async def test_the_stream_is_disabled_and_the_sink_stopped_on_exit(harness, tmp_path): + events, procs, send = harness + sink = video.snapshot(tmp_path / "s.sdp", tmp_path / "f.png") + async with video.StreamSession(sink, send): + pass + assert ("send", video.VIDEO_ENABLE, {"enable": 0}) in events + assert procs[0].terminated + + +async def test_exposure_is_set_then_restored(harness, tmp_path): + events, _, send = harness + sink = video.snapshot(tmp_path / "s.sdp", tmp_path / "f.png") + state = {video.EXPOSURE_STATE_KEY: 0.25} + async with video.StreamSession( + sink, send, exposure=1.0, read_state=lambda keys: {k: state[k] for k in keys} + ): + pass + sends = [e for e in events if e[0] == "send"] + assert sends[0] == ("send", video.EXPOSURE, {"value": 1.0}) + assert sends[-1] == ("send", video.EXPOSURE, {"value": 0.25}) + + +async def test_exposure_outside_the_aircraft_range_is_refused(harness, tmp_path): + _, _, send = harness + with pytest.raises(ValueError): + await video.set_exposure(send, 3.0) + + +async def test_a_refused_videoenable_stops_the_subprocess(harness, tmp_path): + _, procs, _ = harness + sink = video.snapshot(tmp_path / "s.sdp", tmp_path / "f.png") + + async def refusing(name, args): + raise RuntimeError("the drone did not acknowledge VideoEnable") + + with pytest.raises(RuntimeError): + async with video.StreamSession(sink, refusing): + pass + assert procs and procs[0].terminated + + +async def test_a_missing_ffmpeg_is_reported_before_anything_starts(monkeypatch, tmp_path): + monkeypatch.setattr(video.shutil, "which", lambda tool: None) + started = [] + monkeypatch.setattr(video.subprocess, "Popen", lambda *a, **kw: started.append(a)) + + async def send(name, args): + raise AssertionError("must not reach the drone") + + sink = video.snapshot(tmp_path / "s.sdp", tmp_path / "f.png") + with pytest.raises(video.VideoUnavailable): + async with video.StreamSession(sink, send): + pass + assert not started + + +async def test_a_sink_that_outlives_sigterm_is_killed(harness, monkeypatch, tmp_path): + _, procs, send = harness + sink = video.view(tmp_path / "s.sdp") + + class Stubborn(FakeProc): + def terminate(self): + self.terminated = True # but stays running + + def wait(self, timeout=None): + if timeout is not None: + raise subprocess.TimeoutExpired("ffplay", timeout) + return 0 + + def popen(argv, *a, **kw): + procs.append(Stubborn(argv)) + return procs[-1] + + monkeypatch.setattr(video.subprocess, "Popen", popen) + async with video.StreamSession(sink, send): + pass + assert procs[-1].killed + + +async def test_wait_returns_none_while_the_sink_is_still_running(harness, tmp_path): + _, _, send = harness + sink = video.view(tmp_path / "s.sdp") + + session = video.StreamSession(sink, send) + async with session: + session.proc.wait = lambda timeout=None: (_ for _ in ()).throw( + subprocess.TimeoutExpired("ffplay", timeout or 0) + ) + assert await session.wait(0.01) is None + + +def test_the_module_does_not_import_the_parallel_streams(): + # media/ must stay loadable while arsdk/ and protocol/ are still being + # written; it talks to them through a callable the caller supplies. + source = Path(video.__file__).read_text() + assert "import mcbebop.arsdk" not in source + assert "from mcbebop.arsdk" not in source + assert "from mcbebop.protocol" not in source + + +async def test_a_failed_start_puts_the_exposure_back(harness, tmp_path): + # __aexit__ never runs for a failed __aenter__, so the restore has to + # happen on the way out of __aenter__ or the next snapshot inherits it. + events, procs, _ = harness + sink = video.snapshot(tmp_path / "s.sdp", tmp_path / "f.png") + + async def send(name, args): + events.append(("send", name, args)) + if name == video.VIDEO_ENABLE: + raise RuntimeError("the drone did not acknowledge VideoEnable") + + with pytest.raises(RuntimeError): + async with video.StreamSession( + sink, send, exposure=1.0, read_state=lambda keys: {video.EXPOSURE_STATE_KEY: -0.5} + ): + pass + sends = [e for e in events if e[0] == "send"] + assert sends[0] == ("send", video.EXPOSURE, {"value": 1.0}) + assert sends[-1] == ("send", video.EXPOSURE, {"value": -0.5}) + assert procs[0].terminated + + +async def test_a_rejected_exposure_does_not_trip_the_restore_path(harness, tmp_path): + # The exposure was never changed, so there is nothing to put back and no + # NameError on the way out either. + events, procs, send = harness + sink = video.snapshot(tmp_path / "s.sdp", tmp_path / "f.png") + with pytest.raises(ValueError): + async with video.StreamSession(sink, send, exposure=99.0): + pass + assert not [e for e in events if e[0] == "send"] + assert procs[0].terminated