Xianzheng Ma

"""
Unit tests for DecompileAPK extractor, export data, and ClickHouse schemas.

Tests the main DecompileAPK extractor class following the DecompileBinja test pattern.
All analysis is mocked — no JADX/apktool/Java installation required.
"""
import hashlib
import os
import tempfile
import zipfile

import pytest
from unittest.mock import MagicMock, patch, PropertyMock
from dataclasses import asdict

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


def _create_fake_apk_file():
    """Create a minimal APK file for testing and return its path."""
    tmpdir = tempfile.mkdtemp()
    apk_path = os.path.join(tmpdir, "test.apk")
    with zipfile.ZipFile(apk_path, "w") as zf:
        zf.writestr("AndroidManifest.xml", b"<manifest/>")
        zf.writestr("classes.dex", b"dex\n035\x00" + b"\x00" * 100)
    return apk_path


# ============================================================================
# DecompileAPK Extractor Tests
# ============================================================================

class TestDecompileAPK:
    """Tests for the DecompileAPK extractor class."""

    def test_tag(self):
        from redb.extractors.decompiler.DecompileAPK import DecompileAPK
        apk_path = _create_fake_apk_file()
        try:
            ext = DecompileAPK(apk_path, MagicMock())
            assert ext.tag() == "apk_decompiled"
        finally:
            os.unlink(apk_path)

    def test_get_clickhouse_table(self):
        from redb.extractors.decompiler.DecompileAPK import DecompileAPK
        apk_path = _create_fake_apk_file()
        try:
            ext = DecompileAPK(apk_path, MagicMock())
            assert ext.get_clickhouse_table() is None
        finally:
            os.unlink(apk_path)

    def test_timeout_default(self):
        from redb.extractors.decompiler.DecompileAPK import DecompileAPK
        apk_path = _create_fake_apk_file()
        try:
            ext = DecompileAPK(apk_path, MagicMock())
            assert ext.APK_DECOMPILE_TIMEOUT == 600
        finally:
            os.unlink(apk_path)

    def test_timeout_from_env(self):
        from redb.extractors.decompiler.DecompileAPK import DecompileAPK
        apk_path = _create_fake_apk_file()
        try:
            with patch.dict(os.environ, {"APK_DECOMPILE_TIMEOUT": "900"}):
                ext = DecompileAPK(apk_path, MagicMock())
                assert ext.APK_DECOMPILE_TIMEOUT == 900
        finally:
            os.unlink(apk_path)

    def test_timeout_invalid_env(self):
        from redb.extractors.decompiler.DecompileAPK import DecompileAPK
        apk_path = _create_fake_apk_file()
        try:
            with patch.dict(os.environ, {"APK_DECOMPILE_TIMEOUT": "not_a_number"}):
                ext = DecompileAPK(apk_path, MagicMock())
                assert ext.APK_DECOMPILE_TIMEOUT == 600
        finally:
            os.unlink(apk_path)

    def test_calculate_md5(self):
        from redb.extractors.decompiler.DecompileAPK import DecompileAPK
        apk_path = _create_fake_apk_file()
        try:
            ext = DecompileAPK(apk_path, MagicMock())
            result = ext.calculate_md5("test_string")
            assert len(result) == 32
            assert result == hashlib.md5(b"test_string").hexdigest()
        finally:
            os.unlink(apk_path)

    def test_context_manager(self):
        from redb.extractors.decompiler.DecompileAPK import DecompileAPK
        apk_path = _create_fake_apk_file()
        try:
            with DecompileAPK(apk_path, MagicMock()) as ext:
                assert ext is not None
        finally:
            os.unlink(apk_path)

    def test_cleanup_run(self):
        from redb.extractors.decompiler.DecompileAPK import DecompileAPK
        apk_path = _create_fake_apk_file()
        try:
            ext = DecompileAPK(apk_path, MagicMock())
            ext.analyzer = MagicMock()
            ext.cleanup_run()
            assert ext.analyzer is None
        finally:
            os.unlink(apk_path)


# ============================================================================
# Export Data Tests
# ============================================================================

class TestPrepareExportData:
    """Tests for DecompileAPK.prepare_export_data() ClickHouse schemas."""

    def _make_extractor_with_results(self):
        from redb.extractors.decompiler.DecompileAPK import DecompileAPK
        apk_path = _create_fake_apk_file()
        ext = DecompileAPK(apk_path, MagicMock())
        ext.analysis_results = {
            "sha256": "a" * 64,
            "sha1": "b" * 40,
            "md5": "c" * 32,
            "decompiled_content": [
                {
                    "decompiled_method_hash": "d" * 64,
                    "decompiled_method": "public void foo() {}",
                    "decompiled_method_type": "USER",
                    "decompiled_has_string_encryption": False,
                    "decompiled_has_reflection_calls": False,
                    "decompiled_excessive_goto_count": False,
                }
            ],
            "decompiled_refs": [
                {
                    "decompiled_method_hash": "d" * 64,
                    "smali_method_hash": "e" * 64,
                    "decompiled_class_name": "com.example.Foo",
                    "decompiled_method_name": "foo",
                    "decompiled_method_signature": "()V",
                    "decompiled_method_prototype": "void foo()",
                    "functions_caller": ["bar"],
                    "functions_call": ["baz"],
                }
            ],
            "smali_content": [
                {
                    "smali_method_hash": "e" * 64,
                    "smali_method": "invoke-virtual {p0}, Lfoo;->bar()V\nreturn-void",
                    "smali_method_type": "USER",
                    "smali_instructions_count": 10,
                    "smali_register_count": 3,
                    "smali_has_string_encryption": False,
                    "smali_has_reflection_calls": False,
                    "smali_excessive_goto_count": False,
                    "smali_flattened_score": 0.0,
                    "smali_mba_score": 0.0,
                }
            ],
            "smali_refs": [
                {
                    "smali_method_hash": "e" * 64,
                    "decompiled_method_hash": "d" * 64,
                    "smali_class_name": "com.example.Foo",
                    "smali_method_name": "foo",
                    "smali_method_signature": "()V",
                    "ssdeep_smali": None,
                    "tlsh_smali": None,
                }
            ],
            "similarity_metrics": [
                {
                    "smali_method_hash": "e" * 64,
                    "cyclomatic_complexity": 2,
                    "ssdeep_smali": None,
                    "tlsh_smali": None,
                    "ssdeep_smali_normalized": None,
                    "tlsh_smali_normalized": None,
                    "minhash": [],
                }
            ],
            "cfg": [
                {
                    "smali_method_hash": "e" * 64,
                    "cfg_topology_hash": b'\x01' * 16,
                    "block_count": 3,
                    "edge_count": 3,
                    "cfg_instructions_count": 10,
                    "call_count": 1,
                    "cyclomatic_complexity": 2,
                    "loop_count": 0,
                    "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": [[5, 1, 0, 1, 1, 0, 0, 2]] * 3,
                    "cfg_adjacency": [0x00000001, 0x00000002],
                }
            ],
            "strings": [
                {
                    "string": "https://evil.example.com/payload",
                    "string_encoding": "UTF8",
                    "string_offset": 0,
                    "string_length": 31,
                    "string_entropy": 3.95,
                }
            ],
            "analysis_errors": [
                {
                    "class_name": "com.example.Foo",
                    "method_name": "bar",
                    "error_location": "jadx",
                    "error_message": "timeout",
                    "error_type": "TimeoutError",
                }
            ],
        }
        return ext, apk_path

    def test_returns_none_when_no_results(self):
        from redb.extractors.decompiler.DecompileAPK import DecompileAPK
        apk_path = _create_fake_apk_file()
        try:
            ext = DecompileAPK(apk_path, MagicMock())
            assert ext.prepare_export_data("ClickHouseExporter") is None
        finally:
            os.unlink(apk_path)

    def test_multi_table_flag(self):
        ext, apk_path = self._make_extractor_with_results()
        try:
            export = ext.prepare_export_data("ClickHouseExporter")
            assert export["multi_table"] is True
        finally:
            os.unlink(apk_path)

    def test_all_table_keys_present(self):
        ext, apk_path = self._make_extractor_with_results()
        try:
            export = ext.prepare_export_data("ClickHouseExporter")
            assert "decompiled_content" in export
            assert "decompiled_refs" in export
            assert "smali_content" in export
            assert "smali_refs" in export
            assert "method_similarity_metrics" in export
            assert "strings_raw" in export
            assert "analysis_errors" in export
        finally:
            os.unlink(apk_path)

    def test_decompiled_content_schema(self):
        ext, apk_path = self._make_extractor_with_results()
        try:
            export = ext.prepare_export_data("ClickHouseExporter")
            dc = export["decompiled_content"]
            assert dc["table"] == "code_apk_decompiled_methods_content"
            assert len(dc["column_names"]) == 7
            assert len(dc["column_type_names"]) == 7
            assert "decompiled_method_hash" in dc["column_names"]
            assert "decompiled_method" in dc["column_names"]
            assert "analysis_date" in dc["column_names"]
            assert len(dc["data"]) == 1
            assert len(dc["data"][0]) == 7
        finally:
            os.unlink(apk_path)

    def test_decompiled_refs_schema(self):
        ext, apk_path = self._make_extractor_with_results()
        try:
            export = ext.prepare_export_data("ClickHouseExporter")
            dr = export["decompiled_refs"]
            assert dr["table"] == "code_apk_decompiled_methods_references"
            assert len(dr["column_names"]) == 10
            assert len(dr["column_type_names"]) == 10
            assert "sha256" in dr["column_names"]
            assert "sha1" not in dr["column_names"]
            assert "md5" not in dr["column_names"]
            assert "functions_caller" in dr["column_names"]
            assert "functions_call" in dr["column_names"]
        finally:
            os.unlink(apk_path)

    def test_smali_content_schema(self):
        ext, apk_path = self._make_extractor_with_results()
        try:
            export = ext.prepare_export_data("ClickHouseExporter")
            sc = export["smali_content"]
            assert sc["table"] == "code_apk_smali_methods_content"
            assert len(sc["column_names"]) == 11
            assert len(sc["column_type_names"]) == 11
            assert "smali_instructions_count" in sc["column_names"]
            assert "smali_register_count" in sc["column_names"]
            assert "smali_flattened_score" in sc["column_names"]
            assert "smali_mba_score" in sc["column_names"]
        finally:
            os.unlink(apk_path)

    def test_smali_refs_schema(self):
        ext, apk_path = self._make_extractor_with_results()
        try:
            export = ext.prepare_export_data("ClickHouseExporter")
            sr = export["smali_refs"]
            assert sr["table"] == "code_apk_smali_methods_references"
            assert len(sr["column_names"]) == 7
            assert len(sr["column_type_names"]) == 7
            assert "sha1" not in sr["column_names"]
            assert "md5" not in sr["column_names"]
            assert "ssdeep_smali" not in sr["column_names"]
            assert "tlsh_smali" not in sr["column_names"]
        finally:
            os.unlink(apk_path)

    def test_similarity_metrics_schema(self):
        ext, apk_path = self._make_extractor_with_results()
        try:
            export = ext.prepare_export_data("ClickHouseExporter")
            sm = export["method_similarity_metrics"]
            assert sm["table"] == "code_apk_method_similarity_metrics"
            assert len(sm["column_names"]) == 7
            assert "minhash" in sm["column_names"]
            assert "ssdeep_smali_normalized" in sm["column_names"]
            assert "tlsh_smali_normalized" in sm["column_names"]
            # cyclomatic_complexity moved to cfg table
            assert "cyclomatic_complexity" not in sm["column_names"]
            assert "block_count" not in sm["column_names"]
            assert "loop_count" not in sm["column_names"]
        finally:
            os.unlink(apk_path)

    def test_analysis_errors_schema(self):
        ext, apk_path = self._make_extractor_with_results()
        try:
            export = ext.prepare_export_data("ClickHouseExporter")
            ae = export["analysis_errors"]
            assert ae["table"] == "code_apk_analysis_errors"
            assert len(ae["column_names"]) == 9
            assert "error_hash" in ae["column_names"]
            assert "status" in ae["column_names"]
        finally:
            os.unlink(apk_path)

    def test_sha256_propagated_to_refs(self):
        ext, apk_path = self._make_extractor_with_results()
        try:
            export = ext.prepare_export_data("ClickHouseExporter")
            # Check that sha256 from analysis_results appears in ref rows
            dr_row = export["decompiled_refs"]["data"][0]
            assert dr_row[0] == "a" * 64  # sha256
            # sha1 and md5 no longer included in refs
            assert "sha1" not in export["decompiled_refs"]["column_names"]
            assert "md5" not in export["decompiled_refs"]["column_names"]
        finally:
            os.unlink(apk_path)

    def test_empty_tables_not_included(self):
        from redb.extractors.decompiler.DecompileAPK import DecompileAPK
        apk_path = _create_fake_apk_file()
        try:
            ext = DecompileAPK(apk_path, MagicMock())
            ext.analysis_results = {
                "sha256": "a" * 64,
                "sha1": "b" * 40,
                "md5": "c" * 32,
                "decompiled_content": [],
                "decompiled_refs": [],
                "smali_content": [],
                "smali_refs": [],
                "similarity_metrics": [],
                "cfg": [],
                "strings": [],
                "analysis_errors": [],
            }
            export = ext.prepare_export_data("ClickHouseExporter")
            assert export == {"multi_table": True}
        finally:
            os.unlink(apk_path)

    def test_non_clickhouse_exporter_returns_none(self):
        ext, apk_path = self._make_extractor_with_results()
        try:
            result = ext.prepare_export_data("ElasticsearchExporter")
            assert result is None
        finally:
            os.unlink(apk_path)


# ============================================================================
# Dataclass Tests
# ============================================================================

class TestAPKCodeAnalysisDataclasses:
    """Tests for the new APK code analysis dataclasses."""

    def test_decompiled_method_content(self):
        from redb.models.dataclasses import APKDecompiledMethodContent
        dc = APKDecompiledMethodContent(
            decompiled_method_hash="a" * 64,
            decompiled_method="public void foo() {}",
            decompiled_method_type="USER",
        )
        assert dc.decompiled_has_string_encryption is False
        assert dc.decompiled_has_reflection_calls is False
        assert dc.decompiled_excessive_goto_count is False

    def test_decompiled_method_content_populated(self):
        from redb.models.dataclasses import APKDecompiledMethodContent
        dc = APKDecompiledMethodContent(
            decompiled_method_hash="a" * 64,
            decompiled_method="public void foo() {}",
            decompiled_method_type="USER",
            decompiled_has_string_encryption=True,
            decompiled_has_reflection_calls=True,
            decompiled_excessive_goto_count=True,
        )
        assert dc.decompiled_has_string_encryption is True
        d = asdict(dc)
        assert d["decompiled_method_hash"] == "a" * 64

    def test_decompiled_method_reference(self):
        from redb.models.dataclasses import APKDecompiledMethodReference
        dr = APKDecompiledMethodReference(
            sha256="a" * 64,
            sha1="b" * 40,
            md5="c" * 32,
            decompiled_method_hash="d" * 64,
        )
        assert dr.functions_caller == []
        assert dr.functions_call == []
        assert dr.smali_method_hash is None

    def test_smali_method_content(self):
        from redb.models.dataclasses import APKSmaliMethodContent
        sc = APKSmaliMethodContent(
            smali_method_hash="e" * 64,
            smali_method="invoke-virtual {p0}, Lfoo;->bar()V",
            smali_method_type="USER",
            smali_instructions_count=5,
            smali_register_count=3,
        )
        assert sc.smali_instructions_count == 5
        assert sc.smali_register_count == 3

    def test_smali_method_reference(self):
        from redb.models.dataclasses import APKSmaliMethodReference
        sr = APKSmaliMethodReference(
            sha256="a" * 64,
            sha1="b" * 40,
            md5="c" * 32,
            smali_method_hash="e" * 64,
        )
        assert sr.decompiled_method_hash is None
        assert sr.ssdeep_smali is None
        assert sr.tlsh_smali is None

    def test_method_similarity_metrics(self):
        from redb.models.dataclasses import APKMethodSimilarityMetrics
        ms = APKMethodSimilarityMetrics(smali_method_hash="e" * 64)
        assert ms.cyclomatic_complexity is None
        assert ms.ssdeep_smali_normalized is None
        assert ms.tlsh_smali_normalized is None
        assert ms.minhash == []

    def test_code_analysis_error(self):
        from redb.models.dataclasses import APKCodeAnalysisError
        err = APKCodeAnalysisError(
            sha256="a" * 64,
            error_location="jadx",
            error_message="timeout",
        )
        assert err.class_name is None
        assert err.method_name is None
        assert err.error_type is None


# ============================================================================
# Tag Enum Tests
# ============================================================================

class TestAPKDecompiledTag:
    """Tests for the APK_DECOMPILED tag enum."""

    def test_tag_exists(self):
        from redb.extractors.enum import Tag
        assert hasattr(Tag, "APK_DECOMPILED")

    def test_tag_value(self):
        from redb.extractors.enum import Tag
        assert Tag.APK_DECOMPILED.value == "apk_decompiled"


# ============================================================================
# Worker Integration Tests
# ============================================================================

class TestAPKWorkerDecompileIntegration:
    """Tests that DecompileAPK is registered in workers.py."""

    def test_module_registered(self):
        """Verify DecompileAPK is discoverable via the extractor registry."""
        from redb.extractor_registry import get_extractor_class
        cls = get_extractor_class("DecompileAPK")
        assert cls is not None, "DecompileAPK not found in extractor registry"
        assert cls.__name__ == "DecompileAPK"

    def test_worker_apk_dispatch_includes_decompile(self):
        import inspect
        import redb.workers as workers
        source = inspect.getsource(workers.process_binary_file)
        assert "DecompileAPK" in source