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))


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