Ran Zmigrod

"""Mock-based tests for APKCodeAnalyzer._process_method and androguard integration.

Tests the boundary between our code and androguard objects, covering:
- xref extraction (get_xref_from / get_xref_to) for both 4.x and 3.x APIs
- external method skipping
- library method filtering
- smali lookup and androguard fallback disassembly
- deduplication of content hashes
- full output population (all result dict keys)
- minimum instruction filtering
"""
from dataclasses import dataclass
from typing import List
from unittest.mock import MagicMock, Mock, patch, PropertyMock

import pytest

pytestmark = [pytest.mark.unit, pytest.mark.apk, pytest.mark.decompile]


# ---------------------------------------------------------------------------
# Helpers: build mock androguard objects
# ---------------------------------------------------------------------------

SAMPLE_SMALI_BODY = """\
.locals 2
const/4 v0, 0x0
invoke-virtual {p0, v0}, Lcom/example/Foo;->bar(I)V
iget-object v1, p0, Lcom/example/Foo;->name:Ljava/lang/String;
add-int v0, v0, v1
return-void
"""

SAMPLE_JAVA_SOURCE = """\
public void doStuff() {
    int x = 0;
    bar(x);
    String n = this.name;
    x = x + n;
}
"""


def _make_encoded_method(
    class_name="Lcom/example/Foo;",
    method_name="doStuff",
    descriptor="()V",
    code=None,
):
    """Create a mock EncodedMethod."""
    enc = MagicMock()
    enc.get_class_name.return_value = class_name
    enc.get_name.return_value = method_name
    enc.get_descriptor.return_value = descriptor
    enc.get_code.return_value = code
    return enc


def _make_xref_entry_4x(class_name, method_name, offset=0):
    """Create a mock xref tuple in androguard 4.x format.

    4.x: (ClassAnalysis, MethodAnalysis, offset)
    MethodAnalysis has .get_method() returning the EncodedMethod.
    """
    ref_class = MagicMock()
    ref_method = MagicMock()
    inner_enc = MagicMock()
    inner_enc.get_class_name.return_value = class_name
    inner_enc.get_name.return_value = method_name
    ref_method.get_method.return_value = inner_enc
    # Remove direct get_class_name/get_name to simulate 4.x behavior
    del ref_method.get_class_name
    del ref_method.get_name
    return (ref_class, ref_method, offset)


def _make_xref_entry_3x(class_name, method_name, offset=0):
    """Create a mock xref tuple in androguard 3.x format.

    3.x: (ClassAnalysis, EncodedMethod, offset)
    EncodedMethod has .get_class_name() and .get_name() directly.
    """
    ref_class = MagicMock()
    ref_method = MagicMock(spec=[
        "get_class_name", "get_name",
    ])
    ref_method.get_class_name.return_value = class_name
    ref_method.get_name.return_value = method_name
    return (ref_class, ref_method, offset)


def _make_method_analysis(
    encoded=None,
    is_external=False,
    xref_from=None,
    xref_to=None,
):
    """Create a mock MethodAnalysis."""
    method = MagicMock()
    method.is_external.return_value = is_external
    method.get_method.return_value = encoded or _make_encoded_method()
    method.get_xref_from.return_value = xref_from or []
    method.get_xref_to.return_value = xref_to or []
    return method


def _make_analyzer(min_instructions=1):
    """Create an APKCodeAnalyzer with mocked dependencies."""
    with patch(
        "redb.extractors.decompiler.apk.analyzer.JADXDecompiler"
    ), patch(
        "redb.extractors.decompiler.apk.analyzer.ApktoolDisassembler"
    ):
        from redb.extractors.decompiler.apk.analyzer import APKCodeAnalyzer
        analyzer = APKCodeAnalyzer("/fake/path.apk")
        analyzer.min_instructions = min_instructions
        analyzer.library_filter = MagicMock()
        analyzer.library_filter.is_library.return_value = False
        return analyzer


def _empty_results():
    """Create a fresh empty results dict."""
    return {
        "decompiled_content": [],
        "decompiled_refs": [],
        "smali_content": [],
        "smali_refs": [],
        "similarity_metrics": [],
        "cfg": [],
        "analysis_errors": [],
    }


def _make_smali_method(body=SAMPLE_SMALI_BODY, instruction_count=5):
    """Create a mock SmaliMethod."""
    m = MagicMock()
    m.body = body
    m.instruction_count = instruction_count
    m.register_count = 2
    return m


# ---------------------------------------------------------------------------
# Tests: xref extraction
# ---------------------------------------------------------------------------

class TestXrefExtraction:
    """Tests that xrefs are correctly extracted from androguard objects."""

    def test_xref_from_androguard_4x(self):
        """4.x: ref_method.get_method().get_class_name() path."""
        analyzer = _make_analyzer()
        results = _empty_results()
        method_key = "Lcom/example/Foo;->doStuff()V"

        xref_from = [
            _make_xref_entry_4x("Lcom/example/Caller;", "init", 42),
            _make_xref_entry_4x("Lcom/example/Other;", "run", 100),
        ]
        method = _make_method_analysis(xref_from=xref_from)

        smali_methods = {method_key: _make_smali_method()}
        java_methods = {"com.example.Foo.doStuff": SAMPLE_JAVA_SOURCE}

        analyzer._process_method(
            method, smali_methods, java_methods,
            results, set(), set(),
        )

        assert len(results["decompiled_refs"]) == 1
        ref = results["decompiled_refs"][0]
        assert ref["functions_caller"] == [
            "Lcom/example/Caller;->init",
            "Lcom/example/Other;->run",
        ]

    def test_xref_to_androguard_4x(self):
        """4.x: callees extracted via get_method() unwrap."""
        analyzer = _make_analyzer()
        results = _empty_results()
        method_key = "Lcom/example/Foo;->doStuff()V"

        xref_to = [
            _make_xref_entry_4x("Lcom/example/Dep;", "calculate", 10),
        ]
        method = _make_method_analysis(xref_to=xref_to)

        smali_methods = {method_key: _make_smali_method()}
        java_methods = {"com.example.Foo.doStuff": SAMPLE_JAVA_SOURCE}

        analyzer._process_method(
            method, smali_methods, java_methods,
            results, set(), set(),
        )

        ref = results["decompiled_refs"][0]
        assert ref["functions_call"] == ["Lcom/example/Dep;->calculate"]

    def test_xref_from_androguard_3x_fallback(self):
        """3.x: ref_method.get_class_name() direct access fallback."""
        analyzer = _make_analyzer()
        results = _empty_results()
        method_key = "Lcom/example/Foo;->doStuff()V"

        xref_from = [
            _make_xref_entry_3x("Lcom/example/OldCaller;", "legacy"),
        ]
        method = _make_method_analysis(xref_from=xref_from)

        smali_methods = {method_key: _make_smali_method()}
        java_methods = {"com.example.Foo.doStuff": SAMPLE_JAVA_SOURCE}

        analyzer._process_method(
            method, smali_methods, java_methods,
            results, set(), set(),
        )

        ref = results["decompiled_refs"][0]
        assert ref["functions_caller"] == ["Lcom/example/OldCaller;->legacy"]

    def test_xref_to_androguard_3x_fallback(self):
        """3.x: callees via direct access fallback."""
        analyzer = _make_analyzer()
        results = _empty_results()
        method_key = "Lcom/example/Foo;->doStuff()V"

        xref_to = [
            _make_xref_entry_3x("Lcom/example/OldDep;", "compute"),
        ]
        method = _make_method_analysis(xref_to=xref_to)

        smali_methods = {method_key: _make_smali_method()}
        java_methods = {"com.example.Foo.doStuff": SAMPLE_JAVA_SOURCE}

        analyzer._process_method(
            method, smali_methods, java_methods,
            results, set(), set(),
        )

        ref = results["decompiled_refs"][0]
        assert ref["functions_call"] == ["Lcom/example/OldDep;->compute"]

    def test_empty_xrefs(self):
        """No xrefs produces empty lists, not None."""
        analyzer = _make_analyzer()
        results = _empty_results()
        method_key = "Lcom/example/Foo;->doStuff()V"

        method = _make_method_analysis()
        smali_methods = {method_key: _make_smali_method()}
        java_methods = {"com.example.Foo.doStuff": SAMPLE_JAVA_SOURCE}

        analyzer._process_method(
            method, smali_methods, java_methods,
            results, set(), set(),
        )

        ref = results["decompiled_refs"][0]
        assert ref["functions_caller"] == []
        assert ref["functions_call"] == []

    def test_xref_exception_does_not_crash(self):
        """If get_xref_from() itself throws, method still processes."""
        analyzer = _make_analyzer()
        results = _empty_results()
        method_key = "Lcom/example/Foo;->doStuff()V"

        method = _make_method_analysis()
        method.get_xref_from.side_effect = RuntimeError("xref table corrupt")
        method.get_xref_to.side_effect = RuntimeError("xref table corrupt")

        smali_methods = {method_key: _make_smali_method()}
        java_methods = {"com.example.Foo.doStuff": SAMPLE_JAVA_SOURCE}

        analyzer._process_method(
            method, smali_methods, java_methods,
            results, set(), set(),
        )

        # Method should still be processed with empty xrefs
        assert len(results["smali_refs"]) == 1
        ref = results["decompiled_refs"][0]
        assert ref["functions_caller"] == []
        assert ref["functions_call"] == []

    def test_multiple_xrefs(self):
        """Multiple callers and callees are all captured."""
        analyzer = _make_analyzer()
        results = _empty_results()
        method_key = "Lcom/example/Foo;->doStuff()V"

        xref_from = [
            _make_xref_entry_4x("Lcom/A;", "a1"),
            _make_xref_entry_4x("Lcom/B;", "b1"),
            _make_xref_entry_4x("Lcom/C;", "c1"),
        ]
        xref_to = [
            _make_xref_entry_4x("Lcom/D;", "d1"),
            _make_xref_entry_4x("Lcom/E;", "e1"),
        ]
        method = _make_method_analysis(xref_from=xref_from, xref_to=xref_to)

        smali_methods = {method_key: _make_smali_method()}
        java_methods = {"com.example.Foo.doStuff": SAMPLE_JAVA_SOURCE}

        analyzer._process_method(
            method, smali_methods, java_methods,
            results, set(), set(),
        )

        ref = results["decompiled_refs"][0]
        assert len(ref["functions_caller"]) == 3
        assert len(ref["functions_call"]) == 2


# ---------------------------------------------------------------------------
# Tests: method filtering and skipping
# ---------------------------------------------------------------------------

class TestMethodFiltering:
    """Tests for external/library/min-instruction filtering."""

    def test_external_method_skipped(self):
        """External methods (no code body) produce no output."""
        analyzer = _make_analyzer()
        results = _empty_results()

        method = _make_method_analysis(is_external=True)

        analyzer._process_method(
            method, {}, {}, results, set(), set(),
        )

        assert all(len(v) == 0 for v in results.values())

    def test_library_method_skipped(self):
        """Library methods produce no output."""
        analyzer = _make_analyzer()
        analyzer.library_filter.is_library.return_value = True
        results = _empty_results()

        method = _make_method_analysis()

        analyzer._process_method(
            method, {}, {}, results, set(), set(),
        )

        assert all(len(v) == 0 for v in results.values())

    def test_min_instruction_filter(self):
        """Methods below min_instructions threshold are skipped."""
        analyzer = _make_analyzer(min_instructions=10)
        results = _empty_results()
        method_key = "Lcom/example/Foo;->doStuff()V"

        method = _make_method_analysis()
        # Only 3 instructions — below threshold of 10
        smali_methods = {method_key: _make_smali_method(instruction_count=3)}

        analyzer._process_method(
            method, smali_methods, {}, results, set(), set(),
        )

        assert all(len(v) == 0 for v in results.values())

    def test_method_above_min_instructions_processed(self):
        """Methods at or above threshold are processed."""
        analyzer = _make_analyzer(min_instructions=5)
        results = _empty_results()
        method_key = "Lcom/example/Foo;->doStuff()V"

        method = _make_method_analysis()
        smali_methods = {method_key: _make_smali_method(instruction_count=5)}

        analyzer._process_method(
            method, smali_methods, {}, results, set(), set(),
        )

        assert len(results["smali_content"]) == 1
        assert len(results["smali_refs"]) == 1


# ---------------------------------------------------------------------------
# Tests: smali body lookup and androguard fallback
# ---------------------------------------------------------------------------

class TestSmaliLookup:
    """Tests for smali body lookup from apktool and androguard fallback."""

    def test_apktool_smali_used_when_available(self):
        """When apktool smali is available, it's used."""
        analyzer = _make_analyzer()
        results = _empty_results()
        method_key = "Lcom/example/Foo;->doStuff()V"

        method = _make_method_analysis()
        smali_methods = {method_key: _make_smali_method()}

        analyzer._process_method(
            method, smali_methods, {}, results, set(), set(),
        )

        assert len(results["smali_content"]) == 1
        assert results["smali_content"][0]["smali_method"] == SAMPLE_SMALI_BODY

    def test_androguard_fallback_when_no_apktool(self):
        """When apktool smali is missing, _disassemble_with_androguard is used."""
        analyzer = _make_analyzer()
        results = _empty_results()

        # Create encoded method with code
        instruction1 = MagicMock()
        instruction1.get_name.return_value = "const/4"
        instruction1.get_output.return_value = "v0, 0x0"
        instruction2 = MagicMock()
        instruction2.get_name.return_value = "invoke-virtual"
        instruction2.get_output.return_value = "v0, Lcom/Foo;->bar()V"
        instruction3 = MagicMock()
        instruction3.get_name.return_value = "add-int"
        instruction3.get_output.return_value = "v0, v1, v2"
        instruction4 = MagicMock()
        instruction4.get_name.return_value = "sub-int"
        instruction4.get_output.return_value = "v3, v0, v1"
        instruction5 = MagicMock()
        instruction5.get_name.return_value = "return-void"
        instruction5.get_output.return_value = ""

        code = MagicMock()
        code.get_registers_size.return_value = 4
        bytecode = MagicMock()
        bytecode.get_instructions.return_value = [
            instruction1, instruction2, instruction3, instruction4, instruction5,
        ]
        code.get_bc.return_value = bytecode

        encoded = _make_encoded_method(code=code)
        method = _make_method_analysis(encoded=encoded)

        # Empty smali_methods — forces androguard fallback
        analyzer._process_method(
            method, {}, {}, results, set(), set(),
        )

        assert len(results["smali_content"]) == 1
        body = results["smali_content"][0]["smali_method"]
        assert "const/4" in body
        assert "return-void" in body

    def test_androguard_fallback_no_code(self):
        """When encoded method has no code, method is skipped."""
        analyzer = _make_analyzer()
        results = _empty_results()

        encoded = _make_encoded_method(code=None)
        method = _make_method_analysis(encoded=encoded)

        analyzer._process_method(
            method, {}, {}, results, set(), set(),
        )

        assert all(len(v) == 0 for v in results.values())

    def test_androguard_fallback_bytecode_exception(self):
        """When bytecode iteration throws, method is skipped gracefully."""
        analyzer = _make_analyzer()
        results = _empty_results()

        code = MagicMock()
        code.get_registers_size.return_value = 2
        bytecode = MagicMock()
        bytecode.get_instructions.side_effect = RuntimeError("bad bytecode")
        code.get_bc.return_value = bytecode

        encoded = _make_encoded_method(code=code)
        method = _make_method_analysis(encoded=encoded)

        analyzer._process_method(
            method, {}, {}, results, set(), set(),
        )

        # Should not crash, but no content produced
        assert len(results["smali_content"]) == 0


# ---------------------------------------------------------------------------
# Tests: output population and deduplication
# ---------------------------------------------------------------------------

class TestOutputPopulation:
    """Tests that all result dict keys are correctly populated."""

    def _process_standard_method(self, with_java=True):
        """Helper: process a standard method and return results."""
        analyzer = _make_analyzer()
        results = _empty_results()
        method_key = "Lcom/example/Foo;->doStuff()V"

        xref_from = [_make_xref_entry_4x("Lcom/Caller;", "call")]
        xref_to = [_make_xref_entry_4x("Lcom/Dep;", "dep")]
        method = _make_method_analysis(
            xref_from=xref_from, xref_to=xref_to
        )

        smali_methods = {method_key: _make_smali_method()}
        java_methods = {}
        if with_java:
            java_methods = {"com.example.Foo.doStuff": SAMPLE_JAVA_SOURCE}

        analyzer._process_method(
            method, smali_methods, java_methods,
            results, set(), set(),
        )
        return results

    def test_smali_content_populated(self):
        results = self._process_standard_method()
        assert len(results["smali_content"]) == 1
        entry = results["smali_content"][0]
        assert "smali_method_hash" in entry
        assert entry["smali_method"] == SAMPLE_SMALI_BODY
        assert entry["smali_method_type"] == "USER"
        assert entry["smali_instructions_count"] == 5
        assert isinstance(entry["smali_has_string_encryption"], bool)
        assert isinstance(entry["smali_has_reflection_calls"], bool)
        assert isinstance(entry["smali_excessive_goto_count"], bool)

    def test_smali_refs_populated(self):
        results = self._process_standard_method()
        assert len(results["smali_refs"]) == 1
        ref = results["smali_refs"][0]
        assert "smali_method_hash" in ref
        assert ref["smali_class_name"] == "com.example.Foo"
        assert ref["smali_method_name"] == "doStuff"
        assert ref["smali_method_signature"] == "()V"

    def test_decompiled_content_populated_with_java(self):
        results = self._process_standard_method(with_java=True)
        assert len(results["decompiled_content"]) == 1
        entry = results["decompiled_content"][0]
        assert "decompiled_method_hash" in entry
        assert entry["decompiled_method"] == SAMPLE_JAVA_SOURCE

    def test_decompiled_refs_populated_with_java(self):
        results = self._process_standard_method(with_java=True)
        assert len(results["decompiled_refs"]) == 1
        ref = results["decompiled_refs"][0]
        assert "decompiled_method_hash" in ref
        assert "smali_method_hash" in ref
        assert ref["decompiled_class_name"] == "com.example.Foo"
        assert ref["decompiled_method_name"] == "doStuff"
        assert "decompiled_method_prototype" in ref
        assert ref["functions_caller"] == ["Lcom/Caller;->call"]
        assert ref["functions_call"] == ["Lcom/Dep;->dep"]

    def test_no_decompiled_without_java(self):
        results = self._process_standard_method(with_java=False)
        assert len(results["decompiled_content"]) == 0
        assert len(results["decompiled_refs"]) == 0
        # Smali should still be there
        assert len(results["smali_content"]) == 1

    def test_similarity_metrics_populated(self):
        results = self._process_standard_method()
        assert len(results["similarity_metrics"]) == 1
        sim = results["similarity_metrics"][0]
        assert "smali_method_hash" in sim
        assert "cyclomatic_complexity" in sim
        assert "minhash" in sim
        assert isinstance(sim["minhash"], list)
        # block_count moved to cfg table
        assert "block_count" not in sim

    def test_cfg_entry_populated(self):
        results = self._process_standard_method()
        assert len(results["cfg"]) == 1
        cfg = results["cfg"][0]
        assert "smali_method_hash" in cfg
        assert "block_count" in cfg
        assert "cfg_topology_hash" in cfg
        assert "prime_product_smali" in cfg
        assert "wl_minhash" in cfg
        assert "bb_features" in cfg
        assert isinstance(cfg["bb_features"], list)
        for feat in cfg["bb_features"]:
            assert len(feat) == 8

    def test_smali_deduplication(self):
        """Same smali body processed twice produces one content entry."""
        analyzer = _make_analyzer()
        results = _empty_results()
        method_key = "Lcom/example/Foo;->doStuff()V"
        seen_smali = set()

        method = _make_method_analysis()
        smali_methods = {method_key: _make_smali_method()}

        # Process same method twice
        analyzer._process_method(
            method, smali_methods, {}, results, set(), seen_smali,
        )
        analyzer._process_method(
            method, smali_methods, {}, results, set(), seen_smali,
        )

        # Content deduplicated, but refs are per-call
        assert len(results["smali_content"]) == 1
        assert len(results["smali_refs"]) == 2

    def test_decompiled_deduplication(self):
        """Same Java source processed twice produces one content entry."""
        analyzer = _make_analyzer()
        results = _empty_results()
        method_key = "Lcom/example/Foo;->doStuff()V"
        seen_decompiled = set()
        seen_smali = set()

        method = _make_method_analysis()
        smali_methods = {method_key: _make_smali_method()}
        java_methods = {"com.example.Foo.doStuff": SAMPLE_JAVA_SOURCE}

        analyzer._process_method(
            method, smali_methods, java_methods,
            results, seen_decompiled, seen_smali,
        )
        analyzer._process_method(
            method, smali_methods, java_methods,
            results, seen_decompiled, seen_smali,
        )

        assert len(results["decompiled_content"]) == 1
        assert len(results["decompiled_refs"]) == 2


# ---------------------------------------------------------------------------
# Tests: _disassemble_with_androguard
# ---------------------------------------------------------------------------

class TestDisassembleWithAndroguard:
    """Tests for the androguard bytecode disassembly fallback."""

    def test_no_code_returns_none(self):
        from redb.extractors.decompiler.apk.analyzer import APKCodeAnalyzer
        encoded = _make_encoded_method(code=None)
        body, count, regs = APKCodeAnalyzer._disassemble_with_androguard(encoded)
        assert body is None
        assert count == 0
        assert regs == 0

    def test_no_bytecode_returns_none(self):
        from redb.extractors.decompiler.apk.analyzer import APKCodeAnalyzer
        code = MagicMock()
        code.get_registers_size.return_value = 3
        code.get_bc.return_value = None
        encoded = _make_encoded_method(code=code)

        body, count, regs = APKCodeAnalyzer._disassemble_with_androguard(encoded)
        assert body is None
        assert count == 0
        assert regs == 3

    def test_basic_disassembly(self):
        from redb.extractors.decompiler.apk.analyzer import APKCodeAnalyzer

        instructions = []
        for name, output in [
            ("const/4", "v0, 0x0"),
            ("invoke-virtual", "v0, Lcom/Foo;->bar()V"),
            ("return-void", ""),
        ]:
            instr = MagicMock()
            instr.get_name.return_value = name
            instr.get_output.return_value = output
            instructions.append(instr)

        code = MagicMock()
        code.get_registers_size.return_value = 2
        bytecode = MagicMock()
        bytecode.get_instructions.return_value = instructions
        code.get_bc.return_value = bytecode

        encoded = _make_encoded_method(code=code)
        body, count, regs = APKCodeAnalyzer._disassemble_with_androguard(encoded)

        assert body is not None
        assert count == 3
        assert regs == 2
        assert "const/4" in body
        assert "return-void" in body

    def test_invoke_operands_normalized(self):
        """Androguard invoke output gets braces added around registers."""
        from redb.extractors.decompiler.apk.analyzer import APKCodeAnalyzer

        instr = MagicMock()
        instr.get_name.return_value = "invoke-virtual"
        # Androguard format: no braces
        instr.get_output.return_value = "v0, v1, Lcom/Foo;->bar(I)V"

        code = MagicMock()
        code.get_registers_size.return_value = 2
        bytecode = MagicMock()
        bytecode.get_instructions.return_value = [instr]
        code.get_bc.return_value = bytecode

        encoded = _make_encoded_method(code=code)
        body, count, regs = APKCodeAnalyzer._disassemble_with_androguard(encoded)

        # Should be normalized to apktool format with braces
        assert "{v0, v1}" in body

    def test_bytecode_exception_returns_none(self):
        from redb.extractors.decompiler.apk.analyzer import APKCodeAnalyzer

        code = MagicMock()
        code.get_registers_size.return_value = 2
        bytecode = MagicMock()
        bytecode.get_instructions.side_effect = RuntimeError("corrupt")
        code.get_bc.return_value = bytecode

        encoded = _make_encoded_method(code=code)
        body, count, regs = APKCodeAnalyzer._disassemble_with_androguard(encoded)

        assert body is None
        assert regs == 2