#!/usr/bin/env python3
"""
CI tests for guard_cell and QUBO receipt security.

Covers:
- Tamper detection (payload, header)
- Replay denial (cross-container receipt reuse)
- Symlink protection (O_NOFOLLOW)
"""
import os, sys, json, tempfile, struct, base64, hashlib
from pathlib import Path

# Add parent to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))

from tools.guard_cell import guard_open, read_container_facts
from tools.qubo_receipts import make_qubo_bundle
from tools.verify_qubo import verify_qubo

def _make_hx64v2(path, payload=b"test data"):
    """Create minimal hx64v2 container."""
    hdr = {"version": "v0.1.1", "file_size": 0}
    hdr_json = json.dumps(hdr, sort_keys=True, separators=(",",":"))
    hdr_b64 = base64.b64encode(hdr_json.encode()).decode()

    with open(path, "wb") as f:
        f.write(b"hx64v2:")
        f.write(hdr_b64.encode())
        f.write(b"\n")
        f.write(payload)

    # Update file_size in header
    size = os.path.getsize(path)
    hdr["file_size"] = size
    hdr_json = json.dumps(hdr, sort_keys=True, separators=(",",":"))
    hdr_b64 = base64.b64encode(hdr_json.encode()).decode()

    with open(path, "wb") as f:
        f.write(b"hx64v2:")
        f.write(hdr_b64.encode())
        f.write(b"\n")
        f.write(payload)

def _flip_byte(path, offset):
    """Flip one byte at offset."""
    with open(path, "r+b") as f:
        f.seek(offset)
        b = f.read(1)
        if b:
            f.seek(offset)
            f.write(bytes([b[0] ^ 0xFF]))

# --- Tests ---

def test_tamper_payload():
    """Flip byte in payload → guard opens (header-only) but HMAC should fail later."""
    print("🔹 TEST: Tamper payload byte")

    with tempfile.TemporaryDirectory() as td:
        path = Path(td) / "test.hx64v2"
        _make_hx64v2(path, b"clean payload data")

        # Verify opens cleanly first
        try:
            with guard_open(path) as mm:
                assert len(mm) > 0
            print("  ✅ Guard opens clean container")
        except Exception as e:
            print(f"  ❌ Unexpected failure on clean: {e}")
            return False

        # Flip byte in payload (after header)
        _flip_byte(path, os.path.getsize(path) - 5)

        # Guard should still open (header-only checks)
        try:
            with guard_open(path) as mm:
                assert len(mm) > 0
            print("  ✅ Guard opens tampered payload (expected: header-only check)")
        except Exception as e:
            print(f"  ❌ Guard rejected tampered payload (should pass header check): {e}")
            return False

        # NOTE: HMAC/MAC validation would fail here in downstream path
        print("  ℹ️  Downstream HMAC would detect this (not tested here)")

    return True

def test_tamper_header():
    """Flip byte in header → guard_open must fail."""
    print("🔹 TEST: Tamper header byte")

    with tempfile.TemporaryDirectory() as td:
        path = Path(td) / "test.hx64v2"
        _make_hx64v2(path, b"payload")

        # Flip byte in base64 header
        _flip_byte(path, 10)  # In the base64 section

        # Guard should reject
        try:
            with guard_open(path) as mm:
                pass
            print("  ❌ Guard opened tampered header (should fail)")
            return False
        except (ValueError, json.JSONDecodeError, Exception) as e:
            print(f"  ✅ Guard rejected tampered header: {type(e).__name__}")
            return True

def test_replay_denial():
    """Reuse QUBO receipt from container A against container B → verify must fail."""
    print("🔹 TEST: Replay denial (cross-container)")

    with tempfile.TemporaryDirectory() as td:
        path_a = Path(td) / "a.hx64v2"
        path_b = Path(td) / "b.hx64v2"

        _make_hx64v2(path_a, b"payload A")
        _make_hx64v2(path_b, b"payload B different")

        # Create QUBO bundle for A
        bundle_a = make_qubo_bundle(
            str(path_a),
            "test-schema",
            "seed123",
            "FibPi3D-8x8x8",
            [{"params": {"nvars": 64}, "energies": {"e0": -1.5}}]
        )

        lines_a = [ln for ln in bundle_a.strip().split("\n") if ln]

        # Verify against A (should pass)
        result = verify_qubo(lines_a)
        if result["valid"]:
            print("  ✅ Bundle verifies against container A")
        else:
            print(f"  ❌ Bundle failed against A (unexpected): {result['error']}")
            return False

        # Try to use bundle A against container B
        # Extract facts from B
        facts_b = read_container_facts(str(path_b))

        # Modify begin record to point to B (simulating replay attack)
        begin_rec = json.loads(lines_a[0])
        begin_rec["container_sha256"] = facts_b["container_sha256"]
        begin_rec["header_digest"] = facts_b["header_digest"]
        lines_replay = [json.dumps(begin_rec, separators=(",",":"))] + lines_a[1:]

        # Now digest won't match (qubo_end was computed over original begin)
        result = verify_qubo(lines_replay)
        if not result["valid"]:
            print(f"  ✅ Replay attack failed (digest mismatch): {result['error']}")
            return True
        else:
            print("  ❌ Replay attack succeeded (should fail digest check)")
            return False

def test_symlink_trap():
    """Ensure O_NOFOLLOW prevents opening symlinks."""
    print("🔹 TEST: Symlink trap (O_NOFOLLOW)")

    with tempfile.TemporaryDirectory() as td:
        real = Path(td) / "real.hx64v2"
        link = Path(td) / "link.hx64v2"

        _make_hx64v2(real, b"real data")
        os.symlink(real, link)

        # Try to open symlink
        try:
            with guard_open(link) as mm:
                pass
            print("  ❌ Symlink opened (should fail with O_NOFOLLOW)")
            return False
        except (OSError, PermissionError, Exception) as e:
            # Expected: ELOOP (40) or similar
            print(f"  ✅ Symlink rejected: {type(e).__name__}")
            return True

def test_large_file_cap():
    """Test max_bytes size cap."""
    print("🔹 TEST: Large file cap")

    with tempfile.TemporaryDirectory() as td:
        path = Path(td) / "large.hx64v2"
        _make_hx64v2(path, b"data" * 1000)

        # Try to open with tiny cap
        try:
            with guard_open(path, max_bytes=100) as mm:
                pass
            print("  ❌ Large file opened (should fail size cap)")
            return False
        except ValueError as e:
            if "size cap" in str(e):
                print(f"  ✅ Size cap enforced: {e}")
                return True
            else:
                print(f"  ❌ Wrong error: {e}")
                return False

def test_unicode_poison():
    """Inject non-UTF-8 in QUBO receipt → verify must fail."""
    print("🔹 TEST: Unicode poison in QUBO receipt")

    with tempfile.TemporaryDirectory() as td:
        path = Path(td) / "test.hx64v2"
        _make_hx64v2(path, b"payload")

        # Create valid bundle
        bundle = make_qubo_bundle(
            str(path),
            "test-schema",
            "seed123",
            "FibPi3D-8x8x8",
            [{"params": {"nvars": 64}, "energies": {"e0": -1.5}}]
        )

        # Create bytes with invalid UTF-8 then try to encode
        lines = bundle.strip().split("\n")

        # Try to inject invalid UTF-8 - this will fail at encoding
        # So instead, test with valid UTF-8 but malformed JSON
        lines[0] = lines[0].replace("seed123", "seed\u0000")  # Null byte (valid UTF-8, breaks JSON parsing)

        # Verify should fail
        result = verify_qubo(lines)
        if not result["valid"]:
            print(f"  ✅ Malformed input rejected: {result['error']}")
            return True
        else:
            print("  ❌ Malformed input passed (should fail)")
            return False

def main():
    print("━━━ Guard Cell & QUBO Security Tests ━━━\n")

    tests = [
        test_tamper_payload,
        test_tamper_header,
        test_replay_denial,
        test_symlink_trap,
        test_large_file_cap,
        test_unicode_poison,
    ]

    results = []
    for test in tests:
        try:
            passed = test()
            results.append((test.__name__, passed))
        except Exception as e:
            print(f"  ❌ EXCEPTION: {e}")
            results.append((test.__name__, False))
        print()

    print("━━━ Summary ━━━")
    passed = sum(1 for _, p in results if p)
    total = len(results)

    for name, p in results:
        status = "✅" if p else "❌"
        print(f"{status} {name}")

    print(f"\nPassed: {passed}/{total}")

    if passed == total:
        print("✅ All tests passed")
        return 0
    else:
        print("❌ Some tests failed")
        return 1

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