157 lines
6.6 KiB
Python
157 lines
6.6 KiB
Python
"""Voice export/import bundles: roundtrip and hostile-input rejection (temp dirs only)."""
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
import wave
|
|
import zipfile
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
os.environ.setdefault("CONFIG_DIR", tempfile.mkdtemp(prefix="ttvc-bundle-test-config-"))
|
|
|
|
from core import voice_bundle as vb
|
|
from core.voice import _load_meta, _save_meta
|
|
|
|
|
|
def _wav(path: Path):
|
|
with wave.open(str(path), "wb") as w:
|
|
w.setnchannels(1); w.setsampwidth(2); w.setframerate(8000); w.writeframes(b"\x01\x00" * 800)
|
|
|
|
|
|
def _decode(p: Path):
|
|
with wave.open(str(p), "rb") as w:
|
|
assert w.getnframes() > 0
|
|
|
|
|
|
def _zip(files: dict, manifest=None, voice_id="EN_F_a"):
|
|
buf = io.BytesIO()
|
|
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
|
|
for n, d in files.items():
|
|
zf.writestr(n, d)
|
|
if manifest is not False:
|
|
m = manifest or {"format": 1, "voice_id": voice_id,
|
|
"files": {n: hashlib.sha256(d).hexdigest() for n, d in files.items()}}
|
|
zf.writestr("manifest.json", json.dumps(m))
|
|
return buf.getvalue()
|
|
|
|
|
|
class BundleTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = Path(tempfile.mkdtemp(prefix="ttvc-bundle-"))
|
|
self.src, self.dst = self.tmp / "src", self.tmp / "dst"
|
|
self.src.mkdir(); self.dst.mkdir()
|
|
self.wav_bytes = self._wav_bytes()
|
|
|
|
def tearDown(self):
|
|
import shutil; shutil.rmtree(self.tmp, ignore_errors=True)
|
|
|
|
def _wav_bytes(self):
|
|
p = self.tmp / "x.wav"; _wav(p); return p.read_bytes()
|
|
|
|
def _install(self, data: bytes, **kw):
|
|
z = self.tmp / "in.zip"; z.write_bytes(data)
|
|
return vb.import_bundle(z, self.dst, self.dst, decode_check=_decode)
|
|
|
|
def _good(self, **extra):
|
|
files = {"audio.wav": self.wav_bytes, "meta.json": json.dumps({"rating": 4, "seed": 7, **extra}).encode()}
|
|
return files
|
|
|
|
def test_roundtrip_all_parts(self):
|
|
wav = self.src / "EN_F_a.wav"; _wav(wav)
|
|
wav.with_suffix(".reference.txt").write_text("hello")
|
|
wav.with_suffix(".png").write_bytes(b"\x89PNG fake")
|
|
wav.with_suffix(".consent").write_bytes(b"consentaudio")
|
|
_save_meta(wav, {"rating": 3, "voice_design_prompt": "warm", "audio_problems": [],
|
|
"provenance": {"basis": "consented", "source": "https://x.org", "recorded_consent": "EN_F_a.webm"}})
|
|
_, members = vb.export_members(wav)
|
|
data = b"".join(vb.stream_bundle("EN_F_a", members))
|
|
self.assertTrue(zipfile.is_zipfile(io.BytesIO(data)))
|
|
res = self._install(data)
|
|
self.assertEqual(res["status"], "imported")
|
|
out = self.dst / "EN_F_a.wav"
|
|
self.assertEqual(out.read_bytes(), wav.read_bytes())
|
|
self.assertEqual((self.dst / "EN_F_a.reference.txt").read_text(), "hello")
|
|
self.assertEqual((self.dst / "EN_F_a.consent").read_bytes(), b"consentaudio")
|
|
self.assertTrue((self.dst / "EN_F_a.png").exists())
|
|
meta = _load_meta(out)
|
|
self.assertEqual(meta["voice_design_prompt"], "warm")
|
|
self.assertEqual(meta["provenance"]["basis"], "consented")
|
|
self.assertEqual(meta["provenance"]["recorded_consent"], "EN_F_a.webm")
|
|
self.assertEqual([p.name for p in self.dst.iterdir() if p.name.startswith(".import")], [])
|
|
|
|
def test_collision_suffix_and_idempotent(self):
|
|
data = _zip(self._good())
|
|
self.assertEqual(self._install(data)["status"], "imported")
|
|
self.assertEqual(self._install(data)["status"], "unchanged") # same audio again: no duplicate
|
|
self.assertEqual(len(list(self.dst.glob("*.wav"))), 1)
|
|
other = _zip({**self._good(), "audio.wav": self.wav_bytes + b"\x00\x00"})
|
|
r = self._install(other)
|
|
self.assertEqual((r["status"], r["voice_id"]), ("renamed", "EN_F_a_2"))
|
|
self.assertTrue((self.dst / "EN_F_a.wav").exists()) # original untouched
|
|
|
|
def test_missing_optional_parts(self):
|
|
r = self._install(_zip({"audio.wav": self.wav_bytes}))
|
|
self.assertEqual(r["status"], "imported")
|
|
self.assertFalse((self.dst / "EN_F_a.consent").exists())
|
|
|
|
def test_tampered_hash(self):
|
|
files = self._good()
|
|
m = {"format": 1, "voice_id": "EN_F_a", "files": {n: hashlib.sha256(d).hexdigest() for n, d in files.items()}}
|
|
files["audio.wav"] = self.wav_bytes + b"tamper"
|
|
with self.assertRaisesRegex(vb.BundleError, "Hash mismatch"):
|
|
self._install(_zip(files, manifest=m))
|
|
self.assertEqual(list(self.dst.iterdir()), [])
|
|
|
|
def test_zip_slip_and_unknown_members(self):
|
|
for bad in ("../evil.wav", "/etc/passwd", "sub/audio.wav", "run.sh", "audio.exe"):
|
|
with self.assertRaises(vb.BundleError, msg=bad):
|
|
self._install(_zip({**self._good(), bad: b"x"}))
|
|
self.assertEqual(list(self.dst.iterdir()), [])
|
|
|
|
def test_bad_voice_id_in_manifest(self):
|
|
for vid in ("../x", "a/b", ".hidden", "", "a" * 200):
|
|
with self.assertRaises(vb.BundleError, msg=vid):
|
|
self._install(_zip(self._good(), voice_id=vid))
|
|
|
|
def test_zip_bomb(self):
|
|
files = {**self._good(), "reference.txt": b"a" * (vb.MAX_TEXT_BYTES + 1)}
|
|
with self.assertRaises(vb.BundleError):
|
|
self._install(_zip(files))
|
|
files = {"audio.wav": b"\x00" * (20 * 1024 * 1024)} # extreme ratio
|
|
with self.assertRaises(vb.BundleError):
|
|
self._install(_zip(files))
|
|
|
|
def test_lying_declared_size_still_capped(self):
|
|
orig = vb.MAX_TEXT_BYTES
|
|
try:
|
|
data = _zip({**self._good(), "reference.txt": b"b" * 5000})
|
|
vb.MAX_TEXT_BYTES = 2000
|
|
with self.assertRaises(vb.BundleError):
|
|
self._install(data)
|
|
finally:
|
|
vb.MAX_TEXT_BYTES = orig
|
|
|
|
def test_invalid_meta_and_provenance(self):
|
|
for meta in ({"rating": 99}, {"provenance": {"basis": "pirated"}}, {"provenance": {"source": "ftp://u:p@h"}}):
|
|
with self.assertRaises(vb.BundleError):
|
|
self._install(_zip({"audio.wav": self.wav_bytes, "meta.json": json.dumps(meta).encode()}))
|
|
|
|
def test_undecodable_audio_and_not_zip(self):
|
|
with self.assertRaises(vb.BundleError):
|
|
self._install(_zip({"audio.wav": b"not audio at all"}))
|
|
with self.assertRaises(vb.BundleError):
|
|
self._install(b"garbage")
|
|
self.assertEqual(list(self.dst.iterdir()), [])
|
|
|
|
def test_missing_manifest(self):
|
|
with self.assertRaises(vb.BundleError):
|
|
self._install(_zip(self._good(), manifest=False))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|