"""Tests for dataflow analysis (chained reduce calls, suspicious arguments)."""

import pytest

from pickle_scanner import PickleScanner, RiskLevel


class TestDataflowSuspiciousArgs:
    """Test detection of suspicious string arguments to unknown callables."""

    def test_unknown_callable_with_shell_command_arg(self):
        """Unknown callable receiving a shell command should be flagged."""
        # some_module.some_func('echo pwned') — not in threat DB but args
        # contain shell-like content
        malicious = (
            b"csome_module\nsome_func\n"
            b"(S'/bin/sh -c echo pwned'\n"
            b"tR."
        )
        scanner = PickleScanner()
        result = scanner.scan_bytes(malicious, "dataflow_test.pkl")
        # Should flag the suspicious argument
        has_dataflow_finding = any(
            "suspicious" in f.rule.lower() or "dataflow" in f.rule.lower()
            for f in result.report.findings
        )
        assert has_dataflow_finding

    def test_unknown_callable_with_import_arg(self):
        """Unknown callable receiving 'import os' should be flagged."""
        malicious = (
            b"csome_lib\nrun_code\n"
            b"(S'import os; os.system(\"id\")'\n"
            b"tR."
        )
        scanner = PickleScanner()
        result = scanner.scan_bytes(malicious, "import_test.pkl")
        has_dataflow = any(
            "dataflow" in f.rule for f in result.report.findings
        )
        assert has_dataflow

    def test_known_safe_callable_not_flagged_for_args(self):
        """Safe ML pattern callables should NOT trigger dataflow warnings."""
        import pickle
        # A safe pickle with standard dict
        safe_data = pickle.dumps({"key": "import os"})
        scanner = PickleScanner()
        result = scanner.scan_bytes(safe_data, "safe.pkl")
        # No dataflow findings for safe pickles
        dataflow_findings = [
            f for f in result.report.findings
            if "dataflow" in f.rule
        ]
        assert len(dataflow_findings) == 0

    def test_known_dangerous_callable_not_double_flagged(self):
        """Known dangerous callables should already be flagged by threat DB,
        dataflow just adds argument context."""
        malicious = (
            b"cos\nsystem\n"
            b"(S'echo pwned'\n"
            b"tR."
        )
        scanner = PickleScanner()
        result = scanner.scan_bytes(malicious, "os_system.pkl")
        # Should have the primary dangerous_callable finding
        assert result.report.risk_level == RiskLevel.CRITICAL
        # Dataflow suspicious args should NOT trigger for known-dangerous
        # (it only triggers for unknown callables)
        dataflow_suspicious = [
            f for f in result.report.findings
            if f.rule == "dataflow_suspicious_args"
        ]
        assert len(dataflow_suspicious) == 0


class TestDataflowChainedReduce:
    """Test detection of chained reduce call patterns."""

    def test_chained_reduce_detected(self):
        """Two REDUCE calls where one feeds the other should be detected.

        This simulates: os.system(os.popen('id').read())
        which in pickle looks like two chained GLOBAL+REDUCE sequences.
        """
        # First call: os.popen('id') -> REDUCE
        # Second call: os.system(<result>) -> REDUCE
        chained = (
            b"cos\nsystem\n"  # GLOBAL os.system
            b"(cos\npopen\n"  # GLOBAL os.popen (nested)
            b"(S'id'\n"
            b"tR"             # REDUCE os.popen('id')
            b"tR."            # REDUCE os.system(<result>)
        )
        scanner = PickleScanner()
        result = scanner.scan_bytes(chained, "chained.pkl")
        # Should be CRITICAL due to os.system
        assert result.report.risk_level == RiskLevel.CRITICAL


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