Iraj Mahdavi

"""Unit tests for smali semantic normalization.

Tests the three normalization levels (category, opcode, opcode_api),
opcode categorization, ACFG feature index mapping, and integration
with MinHash computation.
"""
import pytest

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


class TestCategorizeOpcode:
    """Tests for opcode -> semantic category mapping."""

    def test_arithmetic(self):
        from redb.extractors.decompiler.apk.smali_normalization import categorize_opcode
        assert categorize_opcode("add-int") == "ALU"
        assert categorize_opcode("mul-long/2addr") == "ALU"
        assert categorize_opcode("neg-int") == "ALU"

    def test_data_movement(self):
        from redb.extractors.decompiler.apk.smali_normalization import categorize_opcode
        assert categorize_opcode("move-result") == "MOV"
        assert categorize_opcode("const/4") == "CONST"
        assert categorize_opcode("const-string") == "CONST"

    def test_memory_access(self):
        from redb.extractors.decompiler.apk.smali_normalization import categorize_opcode
        assert categorize_opcode("iget-object") == "LOAD"
        assert categorize_opcode("sput-wide") == "STORE"
        assert categorize_opcode("aget-byte") == "LOAD"

    def test_invocation(self):
        from redb.extractors.decompiler.apk.smali_normalization import categorize_opcode
        assert categorize_opcode("invoke-virtual") == "CALL"
        assert categorize_opcode("invoke-static") == "CALL"
        assert categorize_opcode("invoke-direct") == "CALL"

    def test_control_flow(self):
        from redb.extractors.decompiler.apk.smali_normalization import categorize_opcode
        assert categorize_opcode("if-eqz") == "BRANCH"
        assert categorize_opcode("goto") == "JMP"
        assert categorize_opcode("goto/16") == "JMP"
        assert categorize_opcode("return-void") == "RET"

    def test_comparison(self):
        from redb.extractors.decompiler.apk.smali_normalization import categorize_opcode
        assert categorize_opcode("cmp-long") == "CMP"
        assert categorize_opcode("cmpl-float") == "CMP"
        assert categorize_opcode("cmpg-double") == "CMP"

    def test_conversion(self):
        from redb.extractors.decompiler.apk.smali_normalization import categorize_opcode
        assert categorize_opcode("int-to-long") == "CONV"
        assert categorize_opcode("float-to-int") == "CONV"

    def test_object_type(self):
        from redb.extractors.decompiler.apk.smali_normalization import categorize_opcode
        assert categorize_opcode("new-instance") == "ALLOC"
        assert categorize_opcode("check-cast") == "TYPE"
        assert categorize_opcode("instance-of") == "TYPE"

    def test_exception_sync(self):
        from redb.extractors.decompiler.apk.smali_normalization import categorize_opcode
        assert categorize_opcode("throw") == "EXC"
        assert categorize_opcode("monitor-enter") == "SYNC"

    def test_unknown(self):
        from redb.extractors.decompiler.apk.smali_normalization import categorize_opcode
        assert categorize_opcode("nop") == "OTHER"


class TestNormalizeInstruction:
    """Tests for single instruction normalization at various levels."""

    def test_category_level(self):
        from redb.extractors.decompiler.apk.smali_normalization import normalize_instruction
        assert normalize_instruction("add-int v0, v1, v2", "category") == "ALU"
        assert normalize_instruction("invoke-virtual {p0}, Lcom/Foo;->bar()V", "category") == "CALL"
        assert normalize_instruction("return-void", "category") == "RET"

    def test_opcode_level_strips_width(self):
        from redb.extractors.decompiler.apk.smali_normalization import normalize_instruction
        # Width variants collapse to same base
        assert normalize_instruction("add-int v0, v1, v2", "opcode") == "add"
        assert normalize_instruction("add-long v0, v2, v4", "opcode") == "add"
        assert normalize_instruction("add-float v0, v1, v2", "opcode") == "add"
        # Addressing modes collapse too
        assert normalize_instruction("add-int/2addr v0, v1", "opcode") == "add"

    def test_opcode_api_invoke(self):
        from redb.extractors.decompiler.apk.smali_normalization import normalize_instruction
        result = normalize_instruction(
            "invoke-virtual {p0, v0}, Lcom/example/Foo;->bar(I)V"
        )
        assert result == "CALL Lcom/example/Foo;->bar(I)V"

    def test_opcode_api_invoke_no_ref(self):
        from redb.extractors.decompiler.apk.smali_normalization import normalize_instruction
        # If somehow no method ref is found, fallback to CALL
        assert normalize_instruction("invoke-virtual {v0}") == "CALL"

    def test_opcode_api_field_load(self):
        from redb.extractors.decompiler.apk.smali_normalization import normalize_instruction
        result = normalize_instruction(
            "iget-object v0, p0, Lcom/Foo;->mField:Ljava/lang/String;"
        )
        assert result == "LOAD Lcom/Foo;->mField:Ljava/lang/String;"

    def test_opcode_api_field_store(self):
        from redb.extractors.decompiler.apk.smali_normalization import normalize_instruction
        result = normalize_instruction(
            "iput-object v0, p0, Lcom/Foo;->mField:Ljava/lang/String;"
        )
        assert result == "STORE Lcom/Foo;->mField:Ljava/lang/String;"

    def test_opcode_api_androguard_field_format(self):
        """androguard uses space instead of colon in field references."""
        from redb.extractors.decompiler.apk.smali_normalization import normalize_instruction
        result = normalize_instruction(
            "iget-object v0, p0, Lcom/Foo;->mField Ljava/lang/String;"
        )
        assert result.startswith("LOAD")
        assert "Lcom/Foo;->mField" in result

    def test_opcode_api_new_instance(self):
        from redb.extractors.decompiler.apk.smali_normalization import normalize_instruction
        result = normalize_instruction(
            "new-instance v0, Lcom/example/MyClass;"
        )
        assert result == "ALLOC Lcom/example/MyClass;"

    def test_opcode_api_const_string(self):
        from redb.extractors.decompiler.apk.smali_normalization import normalize_instruction
        result = normalize_instruction(
            'const-string v0, "hello world"'
        )
        assert result == 'CONST_STR "hello world"'

    def test_opcode_api_const_string_jumbo(self):
        from redb.extractors.decompiler.apk.smali_normalization import normalize_instruction
        result = normalize_instruction(
            'const-string/jumbo v0, "large string"'
        )
        assert result == 'CONST_STR "large string"'

    def test_opcode_api_fallback_to_category(self):
        from redb.extractors.decompiler.apk.smali_normalization import normalize_instruction
        # Non-special instructions fall back to category
        assert normalize_instruction("add-int v0, v1, v2") == "ALU"
        assert normalize_instruction("return-void") == "RET"
        assert normalize_instruction("nop") == "OTHER"

    def test_empty_input(self):
        from redb.extractors.decompiler.apk.smali_normalization import normalize_instruction
        assert normalize_instruction("") == ""
        assert normalize_instruction("   ") == ""


class TestNormalizeMethodBody:
    """Tests for full method body normalization."""

    def test_filters_directives_labels_comments(self):
        from redb.extractors.decompiler.apk.smali_normalization import normalize_method_body
        body = """\
.locals 2
# comment
:label_0
add-int v0, v1, v2
.line 42
invoke-virtual {p0}, Lcom/Foo;->bar()V
return-void
"""
        result = normalize_method_body(body, "category")
        assert result == ["ALU", "CALL", "RET"]

    def test_opcode_api_preserves_api_refs(self):
        from redb.extractors.decompiler.apk.smali_normalization import normalize_method_body
        body = """\
const/4 v0, 0x0
invoke-virtual {p0, v0}, Lcom/Foo;->setX(I)V
iget-object v1, p0, Lcom/Foo;->name:Ljava/lang/String;
return-void
"""
        result = normalize_method_body(body, "opcode_api")
        assert result[0] == "CONST"
        assert result[1] == "CALL Lcom/Foo;->setX(I)V"
        assert result[2] == "LOAD Lcom/Foo;->name:Ljava/lang/String;"
        assert result[3] == "RET"

    def test_empty_body(self):
        from redb.extractors.decompiler.apk.smali_normalization import normalize_method_body
        assert normalize_method_body("") == []
        assert normalize_method_body("\n\n") == []


class TestNormalizationWithMinHash:
    """Integration tests: normalization + MinHash computation."""

    def test_minhash_returns_signature(self):
        from redb.extractors.decompiler.apk.method_extractor import compute_minhash
        body = """\
const/4 v0, 0x0
invoke-virtual {p0, v0}, Lcom/Foo;->bar(I)V
iget-object v1, p0, Lcom/Foo;->name:Ljava/lang/String;
add-int v2, v0, v1
return-void
"""
        sig = compute_minhash(body)
        assert sig is not None
        assert len(sig) == 64
        assert all(0 <= v < 256 for v in sig)

    def test_minhash_register_invariance(self):
        """Same operations with different registers should produce identical signatures."""
        from redb.extractors.decompiler.apk.method_extractor import compute_minhash
        body_a = """\
const/4 v0, 0x0
invoke-virtual {p0, v0}, Lcom/Foo;->bar(I)V
iget-object v1, p0, Lcom/Foo;->name:Ljava/lang/String;
add-int v2, v0, v1
return-void
"""
        body_b = """\
const/4 v3, 0x0
invoke-virtual {p1, v3}, Lcom/Foo;->bar(I)V
iget-object v4, p1, Lcom/Foo;->name:Ljava/lang/String;
add-int v5, v3, v4
return-void
"""
        sig_a = compute_minhash(body_a)
        sig_b = compute_minhash(body_b)
        assert sig_a == sig_b

    def test_minhash_width_invariance(self):
        """Width variants of same operations should produce identical signatures."""
        from redb.extractors.decompiler.apk.method_extractor import compute_minhash
        body_int = """\
add-int v0, v1, v2
sub-int v3, v0, v1
mul-int v4, v0, v3
return v4
"""
        body_long = """\
add-long v0, v2, v4
sub-long v6, v0, v2
mul-long v8, v0, v6
return-wide v8
"""
        sig_int = compute_minhash(body_int)
        sig_long = compute_minhash(body_long)
        # Both normalize to ALU, ALU, ALU, RET -> identical
        assert sig_int == sig_long

    def test_minhash_different_apis_differ(self):
        """Methods calling different APIs should produce different signatures."""
        from redb.extractors.decompiler.apk.method_extractor import compute_minhash
        body_a = """\
const/4 v0, 0x0
invoke-virtual {p0, v0}, Lcom/crypto/AES;->encrypt([B)[B
invoke-virtual {v0}, Ljava/io/OutputStream;->write([B)V
return-void
"""
        body_b = """\
const/4 v0, 0x0
invoke-virtual {p0, v0}, Ljava/lang/String;->length()I
invoke-virtual {v0}, Ljava/io/PrintStream;->println(I)V
return-void
"""
        sig_a = compute_minhash(body_a)
        sig_b = compute_minhash(body_b)
        assert sig_a != sig_b

    def test_minhash_too_few_instructions(self):
        from redb.extractors.decompiler.apk.method_extractor import compute_minhash
        assert compute_minhash("return-void\nnop") is None


class TestACFGBlockFeatures:
    """Tests for per-block ACFG feature extraction via smali_cfg."""

    def test_simple_linear_method(self):
        from redb.extractors.decompiler.apk.smali_cfg import compute_cfg_metrics
        body = """\
add-int v0, v1, v2
invoke-virtual {p0}, Lcom/Foo;->bar()V
return-void
"""
        metrics = compute_cfg_metrics(body)
        assert metrics.block_count >= 1
        assert len(metrics.block_features) == metrics.block_count
        # Each feature vector has 8 elements
        for feat in metrics.block_features:
            assert len(feat) == 8

    def test_branching_method_has_multiple_blocks(self):
        from redb.extractors.decompiler.apk.smali_cfg import compute_cfg_metrics
        body = """\
const/4 v0, 0x0
if-eqz v0, :cond_0
invoke-virtual {p0}, Lcom/Foo;->a()V
goto :goto_0
:cond_0
invoke-virtual {p0}, Lcom/Foo;->b()V
:goto_0
return-void
"""
        metrics = compute_cfg_metrics(body)
        assert metrics.block_count >= 3
        assert len(metrics.block_features) == metrics.block_count

    def test_feature_vector_arithmetic_count(self):
        from redb.extractors.decompiler.apk.smali_cfg import compute_cfg_metrics
        body = """\
add-int v0, v1, v2
sub-int v3, v0, v1
mul-int v4, v0, v3
return v4
"""
        metrics = compute_cfg_metrics(body)
        assert len(metrics.block_features) == 1
        feat = metrics.block_features[0]
        # feat[0] = instr_count (4), feat[1] = arithmetic (3)
        assert feat[0] == 4
        assert feat[1] >= 3  # 3 ALU ops

    def test_feature_vector_call_count(self):
        from redb.extractors.decompiler.apk.smali_cfg import compute_cfg_metrics
        body = """\
invoke-virtual {p0}, Lcom/A;->a()V
invoke-static {v0}, Lcom/B;->b()V
return-void
"""
        metrics = compute_cfg_metrics(body)
        feat = metrics.block_features[0]
        # feat[4] = call count
        assert feat[4] == 2