"""Tests for PickleCloak paper gadget detection.

Verifies that all gadget categories from Table VI of
"The Art of Hide and Seek" (arXiv:2508.19774) are detected.
"""

import pytest

from pickle_scanner import PickleScanner, RiskLevel
from pickle_scanner.analysis.threat_intel import (
    ThreatDatabase,
    is_dangerous_callable,
    RiskLevel as ThreatRiskLevel,
    CRITICAL_CALLABLES,
)


class TestACEGadgets:
    """Test Arbitrary Code Execution gadget detection."""

    @pytest.mark.parametrize("module,name", [
        ("cProfile", "run"),
        ("cProfile", "runctx"),
        ("cProfile.Profile", "run"),
        ("cProfile.Profile", "runctx"),
        ("cgitb", "lookup"),
        ("code.InteractiveInterpreter", "runcode"),
        ("dataclasses", "_create_fn"),
        ("distutils.spawn", "spawn"),
        ("doctest", "debug_script"),
        ("doctest", "_normalize_module"),
        ("profile", "run"),
        ("profile", "runctx"),
        ("profile.Profile", "run"),
        ("profile.Profile", "runctx"),
        ("logging.config", "_resolve"),
        ("logging.config", "_install_handlers"),
        ("pydoc", "pipepager"),
        ("pydoc", "importfile"),
        ("pydoc", "locate"),
        ("pydoc", "safeimport"),
        ("pydoc", "tempfilepager"),
        ("trace.Trace", "run"),
        ("trace.Trace", "runctx"),
        ("unittest.mock", "_importer"),
        ("uuid", "_get_command_stdout"),
        ("_osx_support", "_read_output"),
        ("asyncio.unix_events", "_UnixSubprocessTransport"),
    ])
    def test_stdlib_ace_gadget_detected(self, module, name):
        """Standard library ACE gadgets should be flagged as dangerous."""
        is_dangerous, risk_info = is_dangerous_callable(module, name)
        assert is_dangerous, f"{module}.{name} should be flagged as dangerous"
        assert risk_info.risk_level >= ThreatRiskLevel.CRITICAL

    @pytest.mark.parametrize("module,name", [
        ("numpy.f2py.capi_maps", "getinit"),
        ("numpy.f2py.crackfortran", "_eval_scalar"),
        ("numpy.f2py.crackfortran", "myeval"),
        ("numpy.f2py.crackfortran", "param_eval"),
        ("numpy.f2py.diagnose", "run_command"),
        ("numpy.distutils.cpuinfo", "getoutput"),
        ("numpy.distutils.exec_command", "_exec_command"),
        ("numpy.core.fromnumeric", "_wrapfunc"),
        ("numpy.testing._private.utils", "measure"),
        ("numpy.testing._private.utils", "runstring"),
    ])
    def test_numpy_ace_gadget_detected(self, module, name):
        """NumPy ACE gadgets should be flagged as dangerous."""
        is_dangerous, risk_info = is_dangerous_callable(module, name)
        assert is_dangerous, f"{module}.{name} should be flagged as dangerous"
        assert risk_info.risk_level >= ThreatRiskLevel.CRITICAL

    @pytest.mark.parametrize("module,name", [
        ("sympy", "sympify"),
        ("sympy", "lazy_function"),
        ("sympy.parsing.sympy_parser", "eval_expr"),
        ("sympy.utilities.lambdify", "lambdify"),
        ("sympy.utilities._compilation.compilation", "compile_run_strings"),
        ("sympy.external.importtools", "import_module"),
        ("sympy.polys.monomials.MonomialOps", "_build"),
    ])
    def test_sympy_ace_gadget_detected(self, module, name):
        """SymPy ACE gadgets should be flagged as dangerous."""
        is_dangerous, risk_info = is_dangerous_callable(module, name)
        assert is_dangerous, f"{module}.{name} should be flagged as dangerous"
        assert risk_info.risk_level >= ThreatRiskLevel.CRITICAL

    def test_pandas_ace_gadget_detected(self):
        """Pandas ACE gadgets should be flagged."""
        is_dangerous, risk_info = is_dangerous_callable(
            "pandas._version", "run_command"
        )
        assert is_dangerous
        assert risk_info.risk_level >= ThreatRiskLevel.CRITICAL


class TestFileWriteGadgets:
    """Test File Write gadget detection."""

    @pytest.mark.parametrize("module,name", [
        ("uu", "decode"),
        ("trace.CoverageResults", "write_results_file"),
        ("tracemalloc.Snapshot", "dump"),
        ("profile.Profile", "dump_stats"),
        ("pipes.Template", "open_w"),
        ("mailbox", "_create_carefully"),
        ("http.cookiejar.LWPCookieJar", "save"),
        ("distutils.file_util", "_copy_file_contents"),
        ("distutils.file_util", "write_file"),
        ("xml.etree.ElementTree.ElementTree", "write"),
        ("numpy", "save"),
        ("numpy", "savez"),
        ("numpy", "savez_compressed"),
        ("numpy", "savetxt"),
    ])
    def test_file_write_gadget_detected(self, module, name):
        """File write gadgets should be flagged as HIGH risk."""
        is_dangerous, risk_info = is_dangerous_callable(module, name)
        assert is_dangerous, f"{module}.{name} should be flagged as dangerous"
        assert risk_info.risk_level >= ThreatRiskLevel.HIGH


class TestFileReadGadgets:
    """Test File Read gadget detection."""

    @pytest.mark.parametrize("module,name", [
        ("_pyio", "_open_code_with_warning"),
        ("turtle", "config_dict"),
        ("doctest", "_load_testfile"),
        ("pkgutil.ImpLoader", "get_data"),
        ("pydoc", "_url_handler"),
        ("pipes.Template", "open_r"),
        ("shlex.shlex", "sourcehook"),
        ("argparse.FileType", "__call__"),
        ("fileinput", "hook_compressed"),
        ("numpy", "memmap"),
        ("numpy", "fromfile"),
        ("numpy", "genfromtxt"),
        ("numpy", "loadtxt"),
        ("numpy.DataSource", "open"),
        ("pandas", "read_csv"),
        ("pandas", "read_excel"),
        ("pandas", "read_json"),
    ])
    def test_file_read_gadget_detected(self, module, name):
        """File read gadgets should be flagged as HIGH risk."""
        is_dangerous, risk_info = is_dangerous_callable(module, name)
        assert is_dangerous, f"{module}.{name} should be flagged as dangerous"
        assert risk_info.risk_level >= ThreatRiskLevel.HIGH


class TestNetworkGadgets:
    """Test Network gadget detection."""

    @pytest.mark.parametrize("module,name", [
        ("urllib.request.URLopener", "retrieve"),
        ("urllib.request.URLopener", "open"),
        ("pandas", "read_pickle"),
    ])
    def test_network_gadget_detected(self, module, name):
        """Network gadgets should be flagged."""
        is_dangerous, risk_info = is_dangerous_callable(module, name)
        assert is_dangerous, f"{module}.{name} should be flagged as dangerous"
        assert risk_info.risk_level >= ThreatRiskLevel.HIGH


class TestHelperGadgets:
    """Test helper gadgets (getattr equivalents)."""

    @pytest.mark.parametrize("module,name", [
        ("xmlrpc.server", "resolve_dotted_attribute"),
        ("unittest.mock", "_dot_lookup"),
        ("lib2to3.fixer_util", "attr_chain"),
        ("test.support", "get_attribute"),
    ])
    def test_helper_gadget_detected(self, module, name):
        """Helper gadgets should be flagged at MEDIUM or higher."""
        is_dangerous, risk_info = is_dangerous_callable(module, name)
        assert is_dangerous, f"{module}.{name} should be flagged as dangerous"
        assert risk_info.risk_level >= ThreatRiskLevel.MEDIUM


class TestGadgetPickleDetection:
    """Test that gadgets are detected when embedded in actual pickle bytecode."""

    def test_cprofile_run_pickle_detected(self):
        """A pickle using cProfile.run should be flagged CRITICAL."""
        # Protocol 0 pickle: GLOBAL cProfile run, REDUCE
        malicious = (
            b"ccProfile\nrun\n"
            b"(S'import os; os.system(\"id\")'\n"
            b"tR."
        )
        scanner = PickleScanner()
        result = scanner.scan_bytes(malicious, "cprofile_attack.pkl")
        assert result.report.risk_level == RiskLevel.CRITICAL
        assert any("cProfile" in str(f.message) or "cProfile" in str(f.callable)
                    for f in result.report.findings)

    def test_pydoc_pipepager_pickle_detected(self):
        """A pickle using pydoc.pipepager should be flagged CRITICAL."""
        malicious = (
            b"cpydoc\npipepager\n"
            b"(S'data'\nS'/bin/sh'\n"
            b"tR."
        )
        scanner = PickleScanner()
        result = scanner.scan_bytes(malicious, "pydoc_attack.pkl")
        assert result.report.risk_level == RiskLevel.CRITICAL

    def test_sympy_sympify_pickle_detected(self):
        """A pickle using sympy.sympify should be flagged CRITICAL."""
        malicious = (
            b"csympy\nsympify\n"
            b"(S'__import__(\"os\").system(\"id\")'\n"
            b"tR."
        )
        scanner = PickleScanner()
        result = scanner.scan_bytes(malicious, "sympy_attack.pkl")
        assert result.report.risk_level == RiskLevel.CRITICAL


class TestGadgetCoverage:
    """Test that the gadget database has comprehensive coverage."""

    def test_minimum_gadget_count(self):
        """Database should have at least 130 additional gadgets beyond originals."""
        # Count entries that belong to paper categories
        paper_categories = {
            'code_execution', 'command_execution', 'file_write',
            'file_read', 'introspection', 'deserialization',
        }
        paper_modules = {
            'cProfile', 'cProfile.Profile', 'cgitb', 'code.InteractiveInterpreter',
            'dataclasses', 'distutils.spawn', 'doctest', 'idlelib.calltip',
            'idlelib.run', 'idlelib.run.Executive', 'profile', 'profile.Profile',
            'logging.config', 'pydoc', 'trace', 'trace.Trace',
            'unittest.loader', 'unittest.loader.TestLoader', 'unittest.mock',
            'uuid', '_osx_support', 'asyncio.unix_events',
            'lib2to3.pgen2.grammar', 'lib2to3.pgen2.grammar.Grammar',
            'test.support', 'test.support.PythonSymlink',
            'idlelib.autocomplete', 'idlelib.autocomplete.AutoComplete',
            'idlelib.debugobj', 'idlelib.debugobj.ObjectTreeItem',
            'numpy.f2py.capi_maps', 'numpy.f2py.crackfortran',
            'numpy.f2py.diagnose', 'numpy.distutils.cpuinfo',
            'numpy.distutils.exec_command', 'numpy.distutils.lib2def',
            'numpy.distutils.misc_util', 'numpy.core', 'numpy.core.fromnumeric',
            'numpy.lib.utils', 'numpy.testing._private.utils',
            'numpy.core.tests.test_multiarray',
            'numpy.core.tests.test_multiarray.TestPickling',
            'sympy', 'sympy.parsing.sympy_parser',
            'sympy.physics.mechanics.functions',
            'sympy.printing.tests.test_repr', 'sympy.utilities.source',
            'sympy.utilities.lambdify',
            'sympy.utilities._compilation.compilation',
            'sympy.external.importtools', 'sympy.external.tests.test_codegen',
            'sympy.polys.monomials', 'sympy.polys.monomials.MonomialOps',
            'sympy.plotting.experimental_lambdify',
            'sympy.plotting.experimental_lambdify.Lambdifier',
            'pandas._version', 'uu', 'tracemalloc', 'tracemalloc.Snapshot',
            'pipes', 'pipes.Template', 'mailbox',
            'http.cookiejar', 'http.cookiejar.LWPCookieJar',
            'distutils.file_util', 'xml.etree.ElementTree',
            'xml.etree.ElementTree.ElementTree',
            'sympy.utilities.codegen', 'sympy.utilities.codegen.CodeGen',
            '_pyio', 'turtle', 'pkgutil', 'pkgutil.ImpLoader',
            'shlex', 'shlex.shlex', 'argparse', 'argparse.FileType',
            'fileinput', 'numpy.DataSource',
            'xmlrpc.server', 'lib2to3.fixer_util',
        }
        count = sum(
            1 for (mod, _) in CRITICAL_CALLABLES
            if mod in paper_modules
        )
        # We should have at least 100 entries from the paper
        assert count >= 100, f"Only {count} paper gadgets found in database"


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
