Kaiming Shen

"""
Unit tests for APK CFG feature parity with the Binary Ninja pipeline.

Tests advanced CFG features computed from smali (topology hash, MD-index,
WL-MinHash, prime product, packed adjacency, block features export) and
the slimmed similarity table.

All tests use mocked smali input — no JADX/apktool/Java required.
"""
import pytest

from redb.extractors.decompiler.apk.smali_cfg import (
    SmaliCFGMetrics,
    compute_cfg_metrics,
)
from redb.extractors.decompiler.apk.method_extractor import (
    SMALI_OP_PRIMES,
    compute_prime_product_smali,
    count_call_instructions,
)

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


# ===================================================================
# Sample smali bodies for testing
# ===================================================================

# Simple linear method: 3 instructions, 1 block
SMALI_LINEAR = """\
    const/4 v0, 0x0
    invoke-virtual {p0, v0}, Lcom/Foo;->bar(I)V
    return-void
"""

# Diamond CFG: if-else with two paths merging
SMALI_DIAMOND = """\
    const/4 v0, 0x1
    if-eqz v0, :cond_0
    invoke-virtual {p0}, Lcom/Foo;->pathA()V
    goto :goto_0
    :cond_0
    invoke-virtual {p0}, Lcom/Foo;->pathB()V
    :goto_0
    return-void
"""

# Loop: a back edge from goto to a label before it
SMALI_LOOP = """\
    const/4 v0, 0x0
    :loop_start
    add-int/lit8 v0, v0, 0x1
    if-lt v0, v1, :loop_start
    return-void
"""

# Multi-block with several invoke calls
SMALI_MULTI_CALL = """\
    invoke-virtual {p0}, Lcom/Foo;->a()V
    invoke-static {v0}, Lcom/Bar;->b(I)V
    invoke-direct {p0, v1}, Lcom/Baz;-><init>(I)V
    const/4 v0, 0x0
    return-void
"""


# ===================================================================
# TestSmaliCFGMetricsAdvanced — new fields in SmaliCFGMetrics
# ===================================================================

class TestSmaliCFGMetricsAdvanced:
    """Test that advanced CFG fields are populated by compute_cfg_metrics."""

    def test_topology_hash_is_bytes(self):
        m = compute_cfg_metrics(SMALI_DIAMOND)
        assert isinstance(m.cfg_topology_hash, bytes)
        assert len(m.cfg_topology_hash) == 16

    def test_topology_hash_nonzero_for_multiblock(self):
        m = compute_cfg_metrics(SMALI_DIAMOND)
        assert m.cfg_topology_hash != b'\x00' * 16

    def test_topology_hash_zero_for_empty(self):
        m = compute_cfg_metrics("")
        assert m.cfg_topology_hash == b'\x00' * 16

    def test_identical_cfg_same_topology_hash(self):
        """Two smali methods with the same CFG structure produce the same hash."""
        smali_a = """\
    const/4 v0, 0x1
    if-eqz v0, :cond_0
    const/4 v1, 0x2
    goto :goto_0
    :cond_0
    const/4 v1, 0x3
    :goto_0
    return-void
"""
        smali_b = """\
    const/4 v5, 0x9
    if-nez v5, :cond_0
    add-int v6, v5, v5
    goto :goto_0
    :cond_0
    sub-int v6, v5, v5
    :goto_0
    return-void
"""
        m_a = compute_cfg_metrics(smali_a)
        m_b = compute_cfg_metrics(smali_b)
        assert m_a.cfg_topology_hash == m_b.cfg_topology_hash

    def test_different_cfg_different_topology_hash(self):
        m_linear = compute_cfg_metrics(SMALI_LINEAR)
        m_diamond = compute_cfg_metrics(SMALI_DIAMOND)
        assert m_linear.cfg_topology_hash != m_diamond.cfg_topology_hash

    def test_md_index_topdown_nonzero(self):
        m = compute_cfg_metrics(SMALI_DIAMOND)
        assert m.md_index_topdown != 0

    def test_md_index_bottomup_nonzero(self):
        m = compute_cfg_metrics(SMALI_DIAMOND)
        assert m.md_index_bottomup != 0

    def test_md_index_zero_for_empty(self):
        m = compute_cfg_metrics("")
        assert m.md_index_topdown == 0
        assert m.md_index_bottomup == 0

    def test_wl_minhash_length(self):
        m = compute_cfg_metrics(SMALI_DIAMOND)
        assert len(m.wl_minhash) == 128

    def test_wl_minhash_sentinel_for_empty(self):
        m = compute_cfg_metrics("")
        assert m.wl_minhash == [255] * 128

    def test_wl_minhash_values_are_uint8(self):
        m = compute_cfg_metrics(SMALI_DIAMOND)
        for val in m.wl_minhash:
            assert 0 <= val <= 255

    def test_cfg_adjacency_nonempty_for_multiblock(self):
        m = compute_cfg_metrics(SMALI_DIAMOND)
        assert len(m.cfg_adjacency) > 0

    def test_cfg_adjacency_empty_for_empty(self):
        m = compute_cfg_metrics("")
        assert m.cfg_adjacency == []

    def test_cfg_adjacency_packed_format(self):
        """Each entry should be (src << 16) | tgt."""
        m = compute_cfg_metrics(SMALI_DIAMOND)
        for edge in m.cfg_adjacency:
            assert isinstance(edge, int)
            src = edge >> 16
            tgt = edge & 0xFFFF
            assert src < m.block_count
            assert tgt < m.block_count

    def test_block_features_exported(self):
        m = compute_cfg_metrics(SMALI_DIAMOND)
        assert len(m.block_features) == m.block_count
        for bf in m.block_features:
            assert len(bf) == 8  # 7 categories + successor_count

    def test_cfg_feature_tlsh_none_for_small(self):
        """Small methods produce None (< 50 bytes of feature data)."""
        m = compute_cfg_metrics(SMALI_LINEAR)
        # Single-block linear method — too small for TLSH
        assert m.cfg_feature_tlsh is None

    def test_loop_count_with_back_edge(self):
        m = compute_cfg_metrics(SMALI_LOOP)
        assert m.loop_count >= 1


# ===================================================================
# TestPrimeProductSmali
# ===================================================================

class TestPrimeProductSmali:

    def test_empty_returns_zero(self):
        assert compute_prime_product_smali("") == 0

    def test_none_returns_zero(self):
        assert compute_prime_product_smali(None) == 0

    def test_known_computation(self):
        """Single invoke instruction → CALL category → prime 17."""
        smali = "    invoke-virtual {p0}, Lcom/Foo;->bar()V"
        assert compute_prime_product_smali(smali) == 17

    def test_two_instructions(self):
        """const (CONST→2) * invoke (CALL→17) = 34."""
        smali = """\
    const/4 v0, 0x0
    invoke-virtual {p0, v0}, Lcom/Foo;->bar(I)V
"""
        assert compute_prime_product_smali(smali) == 2 * 17

    def test_skips_directives_and_labels(self):
        """Directives, labels, and comments should be skipped."""
        smali = """\
.registers 2
.line 10
    :label
    # comment
    const/4 v0, 0x0
"""
        # Only const/4 → CONST → prime 2
        assert compute_prime_product_smali(smali) == 2

    def test_position_independent(self):
        """Reordering instructions produces the same product (multiplication is commutative)."""
        smali_a = """\
    add-int v0, v1, v2
    invoke-virtual {p0}, Lcom/Foo;->bar()V
"""
        smali_b = """\
    invoke-virtual {p0}, Lcom/Foo;->bar()V
    add-int v0, v1, v2
"""
        assert compute_prime_product_smali(smali_a) == compute_prime_product_smali(smali_b)

    def test_all_categories_have_primes(self):
        """Every category in SMALI_OP_PRIMES is a positive integer."""
        for cat, prime in SMALI_OP_PRIMES.items():
            assert isinstance(prime, int)
            assert prime >= 1

    def test_mod_2_64(self):
        """Result should be mod 2^64 (doesn't overflow)."""
        # A very long method — product stays within uint64
        smali = "\n".join(["    mul-int v0, v1, v2"] * 1000)
        result = compute_prime_product_smali(smali)
        assert 0 <= result < 2**64


# ===================================================================
# TestCallCount
# ===================================================================

class TestCallCount:

    def test_empty(self):
        assert count_call_instructions("") == 0

    def test_none(self):
        assert count_call_instructions(None) == 0

    def test_single_invoke(self):
        smali = "    invoke-virtual {p0}, Lcom/Foo;->bar()V"
        assert count_call_instructions(smali) == 1

    def test_multiple_invoke_types(self):
        assert count_call_instructions(SMALI_MULTI_CALL) == 3

    def test_non_invoke_not_counted(self):
        smali = """\
    const/4 v0, 0x0
    add-int v1, v0, v0
    return-void
"""
        assert count_call_instructions(smali) == 0


# ===================================================================
# TestSimilarityTableSlimmed
# ===================================================================

class TestSimilarityTableSlimmed:
    """Verify the similarity table no longer contains moved fields."""

    def test_export_similarity_no_block_count(self):
        """block_count should not be in similarity export columns."""
        from redb.extractors.decompiler.DecompileAPK import DecompileAPK
        from unittest.mock import MagicMock

        extractor = DecompileAPK.__new__(DecompileAPK)
        extractor.log = MagicMock()
        extractor.sha256 = "a" * 64
        extractor.sha1 = "b" * 40
        extractor.md5 = "c" * 32
        extractor.analysis_results = {
            "sha256": "a" * 64,
            "sha1": "b" * 40,
            "md5": "c" * 32,
            "similarity_metrics": [{
                "smali_method_hash": "d" * 64,
                "cyclomatic_complexity": 3,
                "ssdeep_smali": None,
                "tlsh_smali": None,
                "minhash": [1, 2, 3],
            }],
        }

        export = extractor.prepare_export_data("ClickHouseExporter")
        sim_cols = export["method_similarity_metrics"]["column_names"]
        assert "block_count" not in sim_cols
        assert "edge_count" not in sim_cols
        assert "loop_count" not in sim_cols
        assert "max_depth" not in sim_cols
        assert "max_fan_out" not in sim_cols
        # These should still be there
        assert "smali_method_hash" in sim_cols
        assert "ssdeep_smali" in sim_cols
        assert "minhash" in sim_cols
        # cyclomatic_complexity moved to cfg table
        assert "cyclomatic_complexity" not in sim_cols


# ===================================================================
# TestCFGTableExport
# ===================================================================

class TestCFGTableExport:
    """Verify the new CFG table is exported correctly."""

    def test_export_cfg_table_present(self):
        from redb.extractors.decompiler.DecompileAPK import DecompileAPK
        from unittest.mock import MagicMock

        extractor = DecompileAPK.__new__(DecompileAPK)
        extractor.log = MagicMock()
        extractor.analysis_results = {
            "sha256": "a" * 64,
            "sha1": "b" * 40,
            "md5": "c" * 32,
            "cfg": [{
                "smali_method_hash": "d" * 64,
                "cfg_topology_hash": b'\x01' * 16,
                "block_count": 4,
                "edge_count": 5,
                "instructions_count": 20,
                "call_count": 3,
                "cyclomatic_complexity": 3,
                "loop_count": 1,
                "max_depth": 2,
                "max_fan_out": 2,
                "md_index_topdown": 12345,
                "md_index_bottomup": 67890,
                "prime_product_smali": 999,
                "cfg_feature_tlsh": None,
                "wl_minhash": [0] * 128,
                "bb_features": [[10, 1, 0, 2, 1, 0, 0, 2]] * 4,
                "cfg_adjacency": [0x00010000, 0x00010002],
            }],
        }

        export = extractor.prepare_export_data("ClickHouseExporter")
        assert "cfg_methods" in export
        cfg_export = export["cfg_methods"]
        assert cfg_export["table"] == "code_apk_cfg_methods"
        assert len(cfg_export["data"]) == 1
        assert "cfg_topology_hash" in cfg_export["column_names"]
        assert "prime_product_smali" in cfg_export["column_names"]
        assert "wl_minhash" in cfg_export["column_names"]
        assert "bb_features" in cfg_export["column_names"]
        assert "cfg_adjacency" in cfg_export["column_names"]
        # Verify column count matches data width
        assert len(cfg_export["column_names"]) == len(cfg_export["data"][0])
        assert len(cfg_export["column_type_names"]) == len(cfg_export["data"][0])