bikini-patchwork/tests/test_abyss.py

from __future__ import annotations

import ast
import base64
import contextlib
import io
import json
import random
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from patchwork import Obfuscator
from patchwork.abyss import AbyssTransformer, build_runtime_stmts, encode_assets


def run_source(source: str) -> str:
    stdout = io.StringIO()
    with contextlib.redirect_stdout(stdout):
        exec(compile(source, "<test>", "exec"), {})
    return stdout.getvalue()


def run_file(path: Path) -> str:
    result = subprocess.run([sys.executable, str(path)], capture_output=True, text=True, timeout=20)
    if result.returncode != 0:
        raise RuntimeError(f"{path} exited {result.returncode}\nstderr:\n{result.stderr}")
    return result.stdout


class AbyssTests(unittest.TestCase):
    def test_transform_removes_protected_body_and_preserves_behavior(self) -> None:
        source = """
def secret(limit):
    marker = "ABYSS_CLEAR_MARKER"
    total = 0
    for item in range(limit):
        if item % 2:
            total += item * 3
        else:
            total += item
    return f"{marker}:{total}"

print(secret(7))
"""
        expected = run_source(source)
        rng = random.Random(2026)
        tree = ast.parse(source)
        transformer = AbyssTransformer(rng, targets={"secret"})
        tree = transformer.protect(tree)
        encoded = encode_assets(transformer.assets, rng)
        tree.body = build_runtime_stmts(encoded, rng) + tree.body
        ast.fix_missing_locations(tree)

        transformed_source = ast.unparse(tree)
        self.assertNotIn("ABYSS_CLEAR_MARKER", transformed_source)
        self.assertIn("secret", transformer.protected)
        self.assertEqual(expected, run_source(transformed_source))

    def test_encoded_assets_have_second_layer_sealing(self) -> None:
        source = """
def secret(limit):
    marker = "ABYSS_SECOND_LAYER_MARKER"
    total = 0
    for item in range(limit):
        total += item
    return f"{marker}:{total}"
"""
        rng = random.Random(7007)
        tree = ast.parse(source)
        transformer = AbyssTransformer(rng, targets={"secret"})
        transformer.protect(tree)
        encoded = encode_assets(transformer.assets, rng)
        outer = base64.b85decode(encoded.payload.encode("ascii"))
        key = base64.b85decode(encoded.key.encode("ascii"))
        decoded = bytes(byte ^ key[index % len(key)] for index, byte in enumerate(outer))
        document = json.loads(decoded.decode("utf-8"))

        self.assertEqual(2, document["v"])
        self.assertIn("p", document)
        self.assertNotIn("ABYSS_SECOND_LAYER_MARKER", decoded.decode("utf-8"))
        self.assertNotIn("CONST", decoded.decode("utf-8"))
        self.assertNotIn("JUMP", decoded.decode("utf-8"))

    def test_obfuscator_abyss_function_can_call_preserved_global(self) -> None:
        source = """
def helper(value):
    return value + 2

def secret(limit):
    total = 0
    for item in range(limit):
        total += helper(item)
    return total

def main():
    print("abyss", secret(6))

if __name__ == "__main__":
    main()
"""
        obfuscator = Obfuscator(seed=99, abyss_functions={"secret"})
        obfuscated = obfuscator.obfuscate(source)
        self.assertEqual(["secret"], obfuscator.abyss_stats["protected"])

        with tempfile.TemporaryDirectory() as raw:
            directory = Path(raw)
            src_path = directory / "sample.py"
            obf_path = directory / "sample_obf.py"
            src_path.write_text(source, encoding="utf-8")
            obf_path.write_text(obfuscated, encoding="utf-8")
            self.assertEqual(run_file(src_path), run_file(obf_path))

    def test_cli_abyss_functions_output_runs(self) -> None:
        source = """
def secret(left, right):
    value = (left << 2) ^ right
    return f"result={value}"

print(secret(9, 5))
"""
        with tempfile.TemporaryDirectory() as raw:
            directory = Path(raw)
            src_path = directory / "sample.py"
            out_path = directory / "sample_obf.py"
            src_path.write_text(source, encoding="utf-8")

            result = subprocess.run(
                [
                    sys.executable,
                    "-m",
                    "patchwork",
                    str(src_path),
                    "-o",
                    str(out_path),
                    "--seed",
                    "17",
                    "--abyss-functions",
                    "secret",
                    "--quiet",
                ],
                cwd=ROOT,
                capture_output=True,
                text=True,
                timeout=20,
            )

            self.assertEqual(0, result.returncode, result.stderr)
            self.assertEqual(run_file(src_path), run_file(out_path))

    def test_public_sample_all_private_functions_are_abyss_supported(self) -> None:
        source = """
from __future__ import annotations
from dataclasses import dataclass
SALT = b"patchwork-disrobe-public-sample-v1"
ALPHABET = "abcdefghijklmnopqrstuvwxyz"
@dataclass(frozen=True)
class Ticket:
    owner: str
    amounts: tuple[int, ...]
def rotate_name(text: str, offset: int) -> str:
    rotated: list[str] = []
    for char in text.lower():
        if char in ALPHABET:
            index = (ALPHABET.index(char) + offset) % len(ALPHABET)
            rotated.append(ALPHABET[index])
        else:
            rotated.append(char)
    return "".join(rotated)
def compact_checksum(ticket: Ticket) -> int:
    state = 0x2A17
    payload = f"{ticket.owner}:{','.join(str(item) for item in ticket.amounts)}".encode()
    for position, byte in enumerate(SALT + payload, start=1):
        state ^= byte * position
        state = ((state << 5) | (state >> 11)) & 0xFFFF
        state = (state + 0x31D + len(ticket.amounts)) & 0xFFFF
    return state
def describe_amounts(amounts: tuple[int, ...]) -> str:
    match amounts:
        case (single,):
            return f"single={single}"
        case (first, second):
            return f"pair={first + second}"
        case (first, *middle, last):
            return f"range={first}->{last};mid={sum(middle)}"
        case _:
            return "empty"
def build_report() -> str:
    tickets = (
        Ticket("Ada", (3, 5, 8, 13)),
        Ticket("Grace", (21, 34)),
        Ticket("Linus", (55,)),
        Ticket("Katherine", (89, 144, 233)),
    )
    lines: list[str] = []
    for index, ticket in enumerate(tickets, start=1):
        alias = rotate_name(ticket.owner, index + 2)
        check = compact_checksum(ticket)
        lines.append(f"{index:02d}|{alias}|{describe_amounts(ticket.amounts)}|{check:04x}")
    return "\\n".join(lines)
if __name__ == "__main__":
    print(build_report())
"""
        obfuscator = Obfuscator(seed=8080, abyss=True)
        obfuscated = obfuscator.obfuscate(source)
        self.assertEqual(["rotate_name", "compact_checksum", "describe_amounts", "build_report"], obfuscator.abyss_stats["protected"])
        self.assertEqual([], obfuscator.abyss_stats["skipped"])
        self.assertNotIn("patchwork-disrobe-public-sample-v1", obfuscated)
        self.assertNotIn("single=", obfuscated)

        with tempfile.TemporaryDirectory() as raw:
            directory = Path(raw)
            src_path = directory / "public_sample.py"
            obf_path = directory / "public_sample_obf.py"
            src_path.write_text(source, encoding="utf-8")
            obf_path.write_text(obfuscated, encoding="utf-8")
            self.assertEqual(run_file(src_path), run_file(obf_path))


if __name__ == "__main__":
    unittest.main()