test_recovery.py

#!/usr/bin/env python3
"""Generate corrupted .torrent variants from good.torrent and test recovery."""

import os
import random
import shutil
import subprocess
import sys
from pathlib import Path

GOOD = Path(__file__).parent / "good.torrent"
TEST_DIR = Path(__file__).parent / "test_corrupted"
SCRIPT = Path(__file__).parent / "torrent_recovery.py"


def setup():
    if TEST_DIR.exists():
        shutil.rmtree(TEST_DIR)
    TEST_DIR.mkdir()


def write(name, data):
    path = TEST_DIR / name
    path.write_bytes(data)
    return path


def make_corrupted_files(data):
    """Generate a variety of corrupted torrent files."""
    files = []

    # 1. Good file (baseline)
    files.append(write("01_good.torrent", data))

    # 2. Truncated at 25%
    files.append(write("02_truncated_25pct.torrent", data[: len(data) // 4]))

    # 3. Truncated at 50%
    files.append(write("03_truncated_50pct.torrent", data[: len(data) // 2]))

    # 4. Truncated at 75%
    files.append(write("04_truncated_75pct.torrent", data[: 3 * len(data) // 4]))

    # 5. Truncated right after the info dict key (before info value)
    marker = data.find(b"4:infod")
    if marker != -1:
        files.append(write("05_truncated_at_info.torrent", data[: marker + 7]))

    # 6. Header only (first 256 bytes)
    files.append(write("06_header_only.torrent", data[:256]))

    # 7. Random bytes sprinkled (corrupt 1% of bytes)
    corrupted = bytearray(data)
    rng = random.Random(42)
    num_flips = max(1, len(data) // 100)
    for _ in range(num_flips):
        pos = rng.randint(0, len(corrupted) - 1)
        corrupted[pos] = rng.randint(0, 255)
    files.append(write("07_random_corruption_1pct.torrent", bytes(corrupted)))

    # 8. Random bytes sprinkled in header area only (first 512 bytes)
    corrupted = bytearray(data)
    for _ in range(50):
        pos = rng.randint(0, min(511, len(corrupted) - 1))
        corrupted[pos] = rng.randint(0, 255)
    files.append(write("08_header_corrupted.torrent", bytes(corrupted)))

    # 9. Garbage prepended (leading junk)
    garbage = bytes(rng.getrandbits(8) for _ in range(128))
    files.append(write("09_garbage_prepended.torrent", garbage + data))

    # 10. Garbage appended
    garbage = bytes(rng.getrandbits(8) for _ in range(128))
    files.append(write("10_garbage_appended.torrent", data + garbage))

    # 11. Middle chunk removed (simulate download gap)
    gap_start = len(data) // 3
    gap_end = gap_start + 2048
    files.append(write("11_middle_gap.torrent", data[:gap_start] + data[gap_end:]))

    # 12. All 'e' terminators replaced with 'x' in first 512 bytes
    corrupted = bytearray(data)
    for i in range(min(512, len(corrupted))):
        if corrupted[i] == ord("e") and i > 0 and (
            corrupted[i - 1] in range(ord("0"), ord("9") + 1)
            or corrupted[i - 1] == ord("e")
        ):
            corrupted[i] = ord("x")
    files.append(write("12_broken_terminators.torrent", bytes(corrupted)))

    # 13. Zero-filled middle section
    corrupted = bytearray(data)
    zero_start = 100
    zero_end = 300
    for i in range(zero_start, min(zero_end, len(corrupted))):
        corrupted[i] = 0
    files.append(write("13_zeroed_header_section.torrent", bytes(corrupted)))

    # 14. Only the info dictionary region
    info_start = data.find(b"4:infod")
    if info_start != -1:
        # Extract a generous chunk from info onward
        files.append(write("14_info_dict_only.torrent", data[info_start:]))

    # 15. Tiny fragment (just 64 bytes)
    files.append(write("15_tiny_fragment.torrent", data[:64]))

    # 16. Empty file
    files.append(write("16_empty.torrent", b""))

    return files


def run_test(path):
    """Run the recovery script on a file, return (success, output)."""
    result = subprocess.run(
        [sys.executable, str(SCRIPT), str(path)],
        capture_output=True,
        text=True,
        timeout=30,
    )
    output = result.stdout + result.stderr
    crashed = result.returncode != 0
    return not crashed, output


def main():
    if not GOOD.exists():
        print(f"Error: {GOOD} not found")
        sys.exit(1)

    data = GOOD.read_bytes()
    print(f"Source: {GOOD.name} ({len(data)} bytes)\n")

    setup()
    files = make_corrupted_files(data)

    passed = 0
    failed = 0
    for path in files:
        name = path.name
        ok, output = run_test(path)
        status = "✅ OK" if ok else "❌ CRASH"
        if not ok:
            failed += 1
        else:
            passed += 1

        # Count what was recovered
        recovered = []
        if "Info Hash" in output and "could not compute" not in output:
            recovered.append("hash")
        if "Name" in output:
            recovered.append("name")
        if "Tracker" in output:
            recovered.append("trackers")
        if "Files" in output or "Single file" in output:
            recovered.append("files")
        if "Piece len" in output:
            recovered.append("pieces")

        tags = ", ".join(recovered) if recovered else "nothing"
        print(f"  {status}  {name:45s}  recovered: {tags}")
        if not ok:
            # Print last 5 lines of error
            lines = output.strip().split("\n")
            for line in lines[-5:]:
                print(f"         {line}")

    print(f"\n{'=' * 60}")
    print(f"  Results: {passed} passed, {failed} crashed out of {len(files)} tests")
    print(f"{'=' * 60}")
    return 0 if failed == 0 else 1


if __name__ == "__main__":
    sys.exit(main())