from __future__ import annotations import hashlib import json import secrets import tempfile import threading import time import unittest from unittest import mock from hub_case import HubCase from hub_engines import log_rows, wav_bytes from hub_support import HOST_STATE, key, signing from kvb_orchestrator import engines, exchange, signed, skill_scripts, voicenotes class TestHubComposer(HubCase): def clip(self, sample: bytes) -> bytes: return wav_bytes(1, 2, engines.STT_RATE, sample * engines.STT_RATE) def record(self, body: bytes, headers: dict[str, str]) -> tuple[int, dict, bytes]: return self.call("POST", "/api/voice", body, {"Content-Type": "audio/wav", **headers}) def note(self, body: bytes) -> str: status, _, answer = self.record(body, self.guard(self.host)) self.assertEqual(status, 202, answer) note = json.loads(answer) self.settle(lambda: self.heard(note["id"])["status"] != "pending") return note["id"] def heard(self, note_id: str) -> dict: status, _, body = self.get(f"/api/voice/{note_id}") self.assertEqual(status, 200, body) return json.loads(body) def submit(self, body: bytes) -> tuple[int, bytes]: status, _, answer = self.record(body, self.guard(self.host)) if status == 202: note = json.loads(answer)["id"] self.settle(lambda: self.heard(note)["status"] != "pending") return status, answer def hear(self, body: bytes) -> tuple[int, bytes]: status, _, answer = self.call("POST", "/api/stt", body, {**self.guard(self.host), "Content-Type": "audio/wav"}) return status, answer def check(self, message_id: str, text: str, to: str) -> dict: return self.tool("operator_message", {"id": message_id, "text": text, "to": to}) def test_a_recording_is_transcribed_off_the_request_and_kept_for_replay(self): good = self.guard(self.host) clip = self.clip(b"\x01\x00") release = threading.Event() whisper = self.httpd.hub.speech.whisper def slow(body: bytes) -> dict: release.wait(10) return whisper(body) with mock.patch.object(self.httpd.hub.speech, "whisper", side_effect=slow): status, _, body = self.record(clip, good) self.assertEqual(status, 202, body) note = json.loads(body) self.assertEqual((note["status"], note["text"], note["seconds"]), ("pending", "", 1.0)) self.assertEqual(self.heard(note["id"])["status"], "pending") release.set() self.settle(lambda: self.heard(note["id"])["status"] == "done") self.assertEqual(self.heard(note["id"])["text"], "hello from the mic second line") status, headers, body = self.get(f"/api/voice/{note['id']}.wav") self.assertEqual((status, headers["content-type"], headers["content-security-policy"], body), (200, "audio/wav", exchange.PLAIN_CSP, clip)) self.assertEqual((self.state / "voice" / f"{note['id']}.wav").read_bytes(), clip) failed = self.note(wav_bytes(1, 2, engines.STT_RATE, b"FAIL" * 100)) self.assertEqual(self.heard(failed)["status"], "failed") self.assertIn("exit code 4", self.heard(failed)["error"]) self.assertEqual(self.get(f"/api/voice/{failed}.wav")[0], 200) for path in ("/api/voice/0123456789abcdef", "/api/voice/0123456789abcdef.wav", "/api/voice/../tts", "/api/voice/x.wav"): with self.subTest(path=path): self.assertEqual(self.get(path)[0], 404) self.assertEqual(self.get("/api/voice")[0], 405) self.assertEqual(len(log_rows(self.engines.scratch / "whisper.jsonl")), 2) self.assertEqual(self.scripts.calls, []) def test_a_recording_is_checked_for_guard_type_size_length_and_load(self): good = self.guard(self.host) clip = self.clip(b"\x01\x00") for headers in ({k: v for k, v in good.items() if k != "X-Kvb-Orchestrator"}, {k: v for k, v in good.items() if k != "Origin"}, {**good, "Origin": "http://evil.example"}): with self.subTest(headers=headers): self.assertEqual(self.record(clip, headers)[0], 403) self.assertEqual(self.call("POST", "/api/voice", clip, {**good, "Content-Type": "audio/webm"})[0], 415) self.assertEqual(self.record(b"RIFF", {**good, "Content-Length": str(engines.STT_LIMIT + 1)})[0], 413) for broken in (b"not a wav", wav_bytes(1, 2, 44100, b"\x00\x00" * 100), wav_bytes(1, 2, engines.STT_RATE, b"")): with self.subTest(broken=broken[:20]): self.assertEqual(self.record(broken, good)[0], 400) with mock.patch.object(engines, "STT_SECONDS", 1): self.assertEqual(self.record(wav_bytes(1, 2, engines.STT_RATE, b"\x00\x00" * (engines.STT_RATE + 1)), good)[0], 400) with mock.patch.object(self.httpd.hub.speech, "input_gap", return_value="the whisper.cpp model m is not a file"): self.assertEqual(self.record(clip, good)[:3:2], (503, b"speech input is off: the whisper.cpp model m is not a file\n")) release = threading.Event() self.addCleanup(release.set) with mock.patch.object(self.httpd.hub.speech, "whisper", side_effect=lambda body: release.wait(10) and {"text": ""}): for _ in range(voicenotes.HEARING_LIMIT): self.assertEqual(self.record(clip, good)[0], 202) self.assertEqual(self.record(clip, good)[0], 503) release.set() self.assertEqual(list((self.state / "voice").glob(".*")), []) def test_only_the_newest_recordings_are_kept(self): with mock.patch.object(voicenotes, "VOICE_KEEP", 2): first, second, third = (self.note(self.clip(bytes([n, 0]))) for n in (1, 2, 3)) self.assertEqual(self.get(f"/api/voice/{first}.wav")[0], 404) self.assertEqual([self.get(f"/api/voice/{note}.wav")[0] for note in (second, third)], [200, 200]) self.assertEqual(sorted(path.stem for path in (self.state / "voice").iterdir()), sorted([second, third])) def test_an_operator_message_needs_the_passkey_and_carries_its_recording(self): self.claim("s1", "Skills session") voice = self.note(self.clip(b"\x01\x00")) payload = {"to": "Skills", "text": " bitte rebasen ", "upload": "", "class": "call", "voice": voice} status, _, body = self.post_json("/api/messages/sign", payload, self.guard(self.host)) self.assertEqual(status, 200, body) self.assertEqual(json.loads(body)["credentials"], []) unpinned = key(-7, "stranger").proof(json.loads(body)["payload"], self.origin) self.assertEqual(self.post_json("/api/messages", {**payload, "proof": unpinned}, self.guard(self.host))[:3:2], (400, b"nicht signiert: no key is pinned\n")) signer = self.pin() proof = self.message_proof(payload, signer) for changed, reason in ( ({"proof": None}, b"nicht signiert: no signature\n"), ({"proof": key(-7, "stranger").proof(proof["payload"], self.origin)}, b"nicht signiert: the signing key is not pinned\n"), ({"proof": signer.proof(proof["payload"], "http://127.0.0.1:1")}, None), ({"text": "bitte mergen", "proof": proof}, b"nicht signiert: the signature does not list this item\n"), ({"class": "post", "proof": proof}, b"nicht signiert: the signature does not list this item\n"), ({"voice": "", "proof": proof}, b"nicht signiert: the signature does not list this item\n"), ): with self.subTest(changed=list(changed)): status, _, answer = self.post_json("/api/messages", {**payload, **changed}, self.guard(self.host)) self.assertEqual(status, 400, answer) if reason is not None: self.assertEqual(answer, reason) self.assertEqual(self.scripts.calls, []) status, _, body = self.post_json("/api/messages", {**payload, "proof": proof}, self.guard(self.host)) self.assertEqual(status, 201, body) sent = json.loads(body) recording = f"{HOST_STATE}/voice/{voice}.wav" self.assertEqual(self.scripts.calls, [(skill_scripts.RADIO, ["send", "call", "--to", "Skills session", "--session", "Operator", f"[kvb-orchestrator msg {sent['id']}] bitte rebasen voice: {recording}"])]) self.assertEqual(self.post_json("/api/messages", {**payload, "proof": proof}, self.guard(self.host))[:3:2], (400, b"nicht signiert: this signature was already used for this item\n")) verified = self.check(sent["id"], "bitte rebasen", "Skills session")["structuredContent"] self.assertEqual((verified["match"], verified["text"], verified["class"], verified["voice"]), (True, "bitte rebasen", "call", recording)) row = self.snapshot()["messages"][-1] self.assertEqual((row["id"], row["class"], row["voice"], row["signed"]), (sent["id"], "call", voice, True)) self.assertEqual(len(self.scripts.calls), 1) def test_a_signed_message_to_all_names_everyone_and_keeps_its_class(self): signer = self.pin() voice = self.note(self.clip(b"\x01\x00")) status, _, body = self.tell_all({"text": "disk full", "class": "mayday", "voice": voice}, signer) self.assertEqual(status, 200, body) script, args = self.scripts.calls[-1] self.assertEqual((script, args[:4]), (skill_scripts.RADIO, ["send", "mayday", "--session", "Operator"])) self.assertTrue(args[-1].endswith(f"] disk full voice: {HOST_STATE}/voice/{voice}.wav"), args[-1]) self.assertEqual(self.snapshot()["messages"][-1]["to"], "*") for payload in ( {"to": "Skills", "text": "x", "upload": "", "class": "mayday", "voice": ""}, {"to": "*", "text": "x", "upload": "0123456789abcdef", "class": "newsletter", "voice": ""}, {"to": "*", "text": "x", "upload": "", "class": "receipt", "voice": ""}, {"to": "*", "text": " ", "upload": "", "class": "broadcast", "voice": ""}, {"to": "*", "text": "x", "upload": "", "class": "broadcast", "voice": "0123456789abcdef"}, {"to": "*", "text": "x", "upload": "", "class": "broadcast", "voice": 7}, {"to": "*", "text": "x", "upload": "", "class": "broadcast"}, ): with self.subTest(payload=payload): self.assertEqual(self.post_json("/api/messages/sign", payload, self.guard(self.host))[0], 400) self.assertEqual(self.post_json("/api/messages/sign", {"to": "*", "text": "x", "upload": "", "class": "broadcast", "voice": ""}, self.client())[0], 403) self.assertEqual(self.get("/api/messages/sign")[0], 405) def test_operator_message_matches_only_the_signed_text_and_recipient(self): self.claim("s1", "Skills session") signer = self.pin() voice = self.note(self.clip(b"\x01\x00")) status, _, body = self.say({"to": "Skills", "text": "bitte\n rebasen", "class": "call", "voice": voice}, signer) self.assertEqual(status, 201, body) sent = json.loads(body)["id"] row = self.scripts.calls[-1][1][-1].split("] ", 1)[1] for text, to in (("bitte rebasen", "Skills"), ("bitte rebasen", "Skills session"), (row, "Skills")): with self.subTest(text=text, to=to): result = self.check(sent, text, to) self.assertFalse(result["isError"], result) found = result["structuredContent"] self.assertEqual((found["match"], found["text"], found["to"], found["class"]), (True, "bitte\n rebasen", "Skills", "call")) self.assertEqual(found["voice"], f"{HOST_STATE}/voice/{voice}.wav") for text, to in (("bitte mergen", "Skills"), ("bitte rebasen", "Stage"), ("bitte rebasen", "*"), ("bitte", "Skills")): with self.subTest(text=text, to=to): result = self.check(sent, text, to) self.assertTrue(result["isError"], result) self.assertNotIn("structuredContent", result) self.assertEqual(self.tell_all({"text": "main is green", "class": "newsletter"}, signer)[0], 200) everyone = next(message["id"] for message in self.snapshot()["messages"] if message["to"] == "*") self.assertTrue(self.check(everyone, "main is green", "*")["structuredContent"]["match"]) self.assertTrue(self.check(everyone, "main is green", "Skills")["isError"]) legacy = self.httpd.hub.messages.add({"to": "Skills", "address": "Skills session", "text": "deploy", "upload": None, "created_ts": time.time()}) self.assertTrue(self.check(legacy["id"], "deploy", "Skills")["isError"]) self.httpd.hub.messages.change(sent, {"text": "rm -rf"}) self.assertTrue(self.check(sent, "rm -rf", "Skills")["isError"]) self.assertTrue(self.tool("operator_message", {"id": everyone})["isError"]) def test_a_post_or_a_call_never_goes_to_everyone(self): signer = self.pin() for cls in ("post", "call"): with self.subTest(cls=cls): payload = {"to": " * ", "text": "x", "upload": "", "class": cls, "voice": ""} self.assertEqual(self.post_json("/api/messages/sign", payload, self.guard(self.host))[:3:2], (400, b"a post or a call names one agent, never *\n")) item = signing.message_item("*", cls, "x", "", "", "") request = signing.approval_payload([signing.item_hash(item)], secrets.token_hex(16), int(time.time()) + 60) proof = signer.proof(request, self.origin) self.assertEqual(self.post_json("/api/messages", {**payload, "proof": proof}, self.guard(self.host))[0], 400) self.assertEqual(self.scripts.calls, []) def test_the_hub_transcribes_a_bounded_number_of_recordings_per_minute(self): answers = [self.submit(self.clip(b"\x01\x00")) for _ in range(12)] statuses = [status for status, _ in answers] self.assertIn(503, statuses) self.assertEqual(statuses, [202] * engines.HEARING_RATE + [503] * (12 - engines.HEARING_RATE)) self.assertIn(f"at most {engines.HEARING_RATE} recordings per {engines.HEARING_WINDOW} s".encode(), answers[-1][1]) self.assertEqual(self.hear(self.clip(b"\x01\x00"))[0], 503) self.assertEqual(len(log_rows(self.engines.scratch / "whisper.jsonl")), engines.HEARING_RATE) def test_the_hub_transcribes_at_most_its_hourly_audio_budget(self): full = wav_bytes(1, 2, engines.STT_RATE, b"\x01\x00" * engines.STT_RATE * engines.STT_SECONDS) statuses = [self.submit(full)[0] for _ in range(5)] status, answer = self.submit(self.clip(b"\x01\x00")) self.assertEqual((statuses, status), ([202] * 5, 503), answer) self.assertEqual(engines.HEARING_BUDGET, 5 * engines.STT_SECONDS) self.assertIn(f"{engines.HEARING_BUDGET} s of audio per hour".encode(), answer) self.assertEqual(self.hear(self.clip(b"\x01\x00"))[0], 503) def test_a_recording_a_stored_message_refers_to_is_never_trimmed(self): self.claim("s1", "Skills session") signer = self.pin() with mock.patch.object(voicenotes, "VOICE_KEEP", 1): kept = self.note(self.clip(b"\x01\x00")) self.assertEqual(self.say({"to": "Skills", "text": "hear this", "class": "post", "voice": kept}, signer)[0], 201) newer = [self.note(self.clip(bytes([n, 0]))) for n in (2, 3)] self.assertEqual(self.get(f"/api/voice/{kept}.wav")[0], 200) self.assertEqual([self.get(f"/api/voice/{note}.wav")[0] for note in newer], [404, 200]) sending = self.note(self.clip(b"\x04\x00")) with self.httpd.hub.voice_notes.held(sending): self.note(self.clip(b"\x05\x00")) self.assertEqual(self.get(f"/api/voice/{sending}.wav")[0], 200) self.assertEqual(self.get(f"/api/voice/{kept}.wav")[0], 200) def test_the_signature_binds_the_bytes_of_the_recording(self): self.claim("s1", "Skills session") signer = self.pin() voice = self.note(self.clip(b"\x01\x00")) wav = self.state / "voice" / f"{voice}.wav" original = wav.read_bytes() payload = {"to": "Skills", "text": "hear this", "upload": "", "class": "call", "voice": voice} proof = self.message_proof(payload, signer) wav.write_bytes(self.clip(b"\x02\x00")) self.assertEqual(self.post_json("/api/messages", {**payload, "proof": proof}, self.guard(self.host))[:3:2], (400, b"nicht signiert: the signature does not list this item\n")) self.assertEqual(self.scripts.calls, []) wav.write_bytes(original) status, _, body = self.post_json("/api/messages", {**payload, "proof": proof}, self.guard(self.host)) self.assertEqual(status, 201, body) found = self.check(json.loads(body)["id"], "hear this", "Skills")["structuredContent"] self.assertEqual(found["voice_sha256"], hashlib.sha256(original).hexdigest()) def test_a_signature_check_without_a_temporary_directory_answers_400(self): self.claim("s1", "Skills session") signer = self.pin() payload = {"to": "Skills", "text": "x", "upload": "", "class": "post", "voice": ""} proof = self.message_proof(payload, signer) missing = self.state / "no-such-tmp" with mock.patch.object(signed.passkey, "serialization", None), mock.patch.object(tempfile, "tempdir", str(missing)): status, _, body = self.post_json("/api/messages", {**payload, "proof": proof}, self.guard(self.host)) self.assertEqual(status, 400, body) self.assertTrue(body.startswith(b"nicht signiert: "), body) self.assertEqual(self.scripts.calls, []) if __name__ == "__main__": unittest.main()