Ran Zhang

"""Unit tests (mocked BN) for bninja/decompiler.py — BinaryNinjaDecompiler."""
import sys
import json
import time
import pytest
from unittest.mock import MagicMock, patch, mock_open

from tests.unit.conftest_binja_stubs import (
    install_binja_stubs,
    MockFunction,
    MockBasicBlock,
    MockDisassemblyLine,
    MockToken,
    MockSymbol,
    MockBinaryView,
    MockEdge,
    SymbolType,
    InstructionTextTokenType,
    BranchType,
)

install_binja_stubs()

from redb.extractors.decompiler.bninja.decompiler import (
    BinaryNinjaDecompiler,
    is_lib_or_thunk,
)
from redb.extractors.decompiler.bninja.function_type import FunctionType


# ============================================================================
# 10a. BinaryNinjaDecompiler
# ============================================================================


class TestBinaryNinjaDecompilerLogError:
    def _make_decompiler(self):
        with patch.object(BinaryNinjaDecompiler, "__init__", lambda self, *a, **kw: None):
            d = BinaryNinjaDecompiler.__new__(BinaryNinjaDecompiler)
            d.log = MagicMock()
            d.errors = []
            d.filepath = "/fake/binary"
            d.bv = None
            d.analysis_results = None
            d.exporters = []
            d.index_prefix = None
            d.filetype = None
            d.goresym = None
            d.BNINJA_TIMEOUT = 60
            return d

    def test_log_error_appends(self):
        d = self._make_decompiler()
        d.log_error("test error", "func1", 0x1000, Exception("boom"), "extract")
        assert len(d.errors) == 1

    def test_log_error_schema(self):
        d = self._make_decompiler()
        d.log_error("test error", "func1", 0x1000, Exception("boom"), "extract")
        error = d.errors[0]
        expected_keys = [
            "function_name",
            "function_address",
            "error_location",
            "error_message",
            "error_details",
            "error_type",
            "timestamp",
        ]
        for key in expected_keys:
            assert key in error, f"Missing key: {key}"


class TestBinaryNinjaDecompilerIsTooFewBlocks:
    def _make_decompiler(self):
        with patch.object(BinaryNinjaDecompiler, "__init__", lambda self, *a, **kw: None):
            d = BinaryNinjaDecompiler.__new__(BinaryNinjaDecompiler)
            d.log = MagicMock()
            d.errors = []
            d.MIN_FUNCTION_SIZE = 10
            return d

    def test_is_too_few_blocks_none(self):
        d = self._make_decompiler()
        func = MagicMock()
        func.basic_blocks = None
        assert d.is_too_few_blocks(func) is False

    def test_is_too_few_blocks_empty(self):
        # Empty basic_blocks now returns True because Binja creates function
        # stubs with empty blocks before analysis completes — we no longer
        # skip these so they get filtered downstream instead.
        d = self._make_decompiler()
        func = MagicMock()
        func.basic_blocks = []
        assert d.is_too_few_blocks(func) is True

    def test_is_too_few_blocks_small_single(self):
        d = self._make_decompiler()
        block = MagicMock()
        block.disassembly_text = [MagicMock() for _ in range(5)]  # < 10
        func = MagicMock()
        func.basic_blocks = [block]
        assert d.is_too_few_blocks(func) is False

    def test_is_too_few_blocks_large_single(self):
        d = self._make_decompiler()
        block = MagicMock()
        block.disassembly_text = [MagicMock() for _ in range(15)]  # >= 10
        func = MagicMock()
        func.basic_blocks = [block]
        assert d.is_too_few_blocks(func) is True

    def test_is_too_few_blocks_multiple(self):
        d = self._make_decompiler()
        func = MagicMock()
        func.basic_blocks = [MagicMock(), MagicMock()]
        assert d.is_too_few_blocks(func) is True


class TestIsLibOrThunk:
    def test_is_lib_or_thunk_user(self):
        func = MagicMock()
        func.symbol.type = SymbolType.FunctionSymbol
        func.is_thunk = False
        assert is_lib_or_thunk(func) is True  # NOT filtered

    def test_is_lib_or_thunk_thunk(self):
        func = MagicMock()
        func.symbol.type = SymbolType.FunctionSymbol
        func.is_thunk = True
        assert is_lib_or_thunk(func) is False  # Filtered out

    def test_is_lib_or_thunk_external(self):
        func = MagicMock()
        func.symbol.type = SymbolType.ImportedFunctionSymbol
        func.is_thunk = False
        assert is_lib_or_thunk(func) is False  # Filtered out


class TestBinaryNinjaDecompilerExtract:
    def _make_decompiler(self):
        with patch.object(BinaryNinjaDecompiler, "__init__", lambda self, *a, **kw: None):
            d = BinaryNinjaDecompiler.__new__(BinaryNinjaDecompiler)
            d.log = MagicMock()
            d.errors = []
            d.filepath = "/fake/binary"
            d.bv = MagicMock()
            d.analysis_results = None
            d.exporters = []
            d.index_prefix = None
            d.filetype = None
            d.goresym = None
            d.BNINJA_TIMEOUT = 60
            d.arch = MagicMock()
            return d

    def test_extract_success(self):
        d = self._make_decompiler()
        mock_results = {
            "decompiled": [{"test": True}],
            "disassembled": [],
            "cfg": [],
            "llil": [],
            "errors": [],
            "strings": [],
            "mlil": [],
        }
        with patch.object(d, "analyze_binary", return_value=mock_results):
            result = d.extract()
            assert result is True
            assert d.analysis_results == mock_results

    def test_extract_failure(self):
        d = self._make_decompiler()
        with patch.object(d, "analyze_binary", return_value=None):
            result = d.extract()
            assert result is False

    def test_tag(self):
        d = self._make_decompiler()
        assert d.tag() == "DECOMPILED"


class TestBinaryNinjaDecompilerAnalyzeBinary:
    def _make_decompiler(self):
        with patch.object(BinaryNinjaDecompiler, "__init__", lambda self, *a, **kw: None):
            d = BinaryNinjaDecompiler.__new__(BinaryNinjaDecompiler)
            d.log = MagicMock()
            d.errors = []
            d.filepath = "/fake/binary"
            d.analysis_results = None
            d.exporters = []
            d.index_prefix = None
            d.filetype = None
            d.goresym = None
            d.BNINJA_TIMEOUT = 60
            d.arch = MagicMock()
            d.MIN_FUNCTION_SIZE = 10
            return d

    def test_analyze_binary_links_hlil_disasm(self):
        """Bi-directional hash linkage between decompiled and disassembled."""
        d = self._make_decompiler()
        d.bv = MagicMock()
        d.bv.functions = []
        # Mock methods to return controlled data
        with patch.object(d, "extract_hlil") as mock_hlil, \
             patch.object(d, "extract_disasm") as mock_disasm, \
             patch.object(d, "extract_cfg", return_value=None), \
             patch.object(d, "extract_lowlevel", return_value=None):

            mock_hlil.return_value = {
                "decompiled_function_hash": "HLIL_HASH",
                "decompiled_function_name": "test",
                "decompiled_function": "code",
            }
            mock_disasm.return_value = {
                "disassembled_function_hash": "DISASM_HASH",
                "disassembled_function_no_addresses": "asm",
            }

            # Manually feed one function
            mock_func = MagicMock()
            mock_func.symbol.type = SymbolType.FunctionSymbol
            mock_func.is_thunk = False
            mock_func.basic_blocks = [MagicMock(), MagicMock()]
            d.bv.functions = [mock_func]

            results = d.analyze_binary()
            if results and results["decompiled"]:
                hlil_r = results["decompiled"][0]
                disasm_r = results["disassembled"][0]
                assert hlil_r["disassembled_function_hash"] == "DISASM_HASH"
                assert disasm_r["decompiled_function_hash"] == "HLIL_HASH"

    def test_analyze_binary_links_llil_disasm(self):
        """Bi-directional linkage between LLIL and disassembly."""
        d = self._make_decompiler()
        d.bv = MagicMock()

        with patch.object(d, "extract_hlil", return_value=None), \
             patch.object(d, "extract_disasm") as mock_disasm, \
             patch.object(d, "extract_cfg", return_value=None), \
             patch.object(d, "extract_lowlevel") as mock_llil:

            mock_disasm.return_value = {
                "disassembled_function_hash": "DISASM_HASH",
                "disassembled_function_no_addresses": "asm",
            }
            mock_llil.return_value = {
                "sha256_llil": "LLIL_HASH",
                "tlsh_llil": "tlsh_val",
            }

            mock_func = MagicMock()
            mock_func.symbol.type = SymbolType.FunctionSymbol
            mock_func.is_thunk = False
            mock_func.basic_blocks = [MagicMock(), MagicMock()]
            d.bv.functions = [mock_func]

            results = d.analyze_binary()
            if results and results["llil"]:
                llil_r = results["llil"][0]
                assert llil_r["disassembled_function_hash"] == "DISASM_HASH"

    def test_analyze_binary_links_cfg_disasm(self):
        """CFG gets disassembled_function_hash."""
        d = self._make_decompiler()
        d.bv = MagicMock()

        with patch.object(d, "extract_hlil", return_value=None), \
             patch.object(d, "extract_disasm") as mock_disasm, \
             patch.object(d, "extract_cfg") as mock_cfg, \
             patch.object(d, "extract_lowlevel", return_value=None):

            mock_disasm.return_value = {
                "disassembled_function_hash": "DISASM_HASH",
                "disassembled_function_no_addresses": "asm",
            }
            mock_cfg.return_value = {
                "function_address": 0x1000,
                "cyclomatic_complexity": 3,
            }

            mock_func = MagicMock()
            mock_func.symbol.type = SymbolType.FunctionSymbol
            mock_func.is_thunk = False
            mock_func.basic_blocks = [MagicMock(), MagicMock()]
            d.bv.functions = [mock_func]

            results = d.analyze_binary()
            if results and results["cfg"]:
                cfg_r = results["cfg"][0]
                assert cfg_r["disassembled_function_hash"] == "DISASM_HASH"

    def test_analyze_binary_sets_cyclomatic(self):
        """cyclomatic_complexity set on disasm from CFG."""
        d = self._make_decompiler()
        d.bv = MagicMock()

        with patch.object(d, "extract_hlil", return_value=None), \
             patch.object(d, "extract_disasm") as mock_disasm, \
             patch.object(d, "extract_cfg") as mock_cfg, \
             patch.object(d, "extract_lowlevel", return_value=None):

            mock_disasm.return_value = {
                "disassembled_function_hash": "HASH",
                "disassembled_function_no_addresses": "asm",
            }
            mock_cfg.return_value = {
                "function_address": 0x1000,
                "cyclomatic_complexity": 7,
            }

            mock_func = MagicMock()
            mock_func.symbol.type = SymbolType.FunctionSymbol
            mock_func.is_thunk = False
            mock_func.basic_blocks = [MagicMock(), MagicMock()]
            d.bv.functions = [mock_func]

            results = d.analyze_binary()
            if results and results["disassembled"]:
                disasm_r = results["disassembled"][0]
                assert disasm_r.get("cyclomatic_complexity") == 7


class TestApplyGoReSym:
    def _make_decompiler(self):
        with patch.object(BinaryNinjaDecompiler, "__init__", lambda self, *a, **kw: None):
            d = BinaryNinjaDecompiler.__new__(BinaryNinjaDecompiler)
            d.log = MagicMock()
            d.errors = []
            d.bv = MagicMock()
            return d

    def test_apply_goresym_user_functions(self):
        d = self._make_decompiler()
        d.goresym = "/tmp/goresym.json"
        data = {
            "UserFunctions": [{"Start": 0x1000, "FullName": "main.main"}],
        }
        with patch("builtins.open", mock_open(read_data=json.dumps(data))):
            d._BinaryNinjaDecompiler__apply_goresym()
            d.bv.define_user_symbol.assert_called()

    def test_apply_goresym_std_functions(self):
        d = self._make_decompiler()
        d.goresym = "/tmp/goresym.json"
        data = {
            "StdFunctions": [{"Start": 0x2000, "FullName": "runtime.main"}],
        }
        with patch("builtins.open", mock_open(read_data=json.dumps(data))):
            d._BinaryNinjaDecompiler__apply_goresym()
            d.bv.define_user_symbol.assert_called()

    def test_apply_goresym_invalid_json(self):
        d = self._make_decompiler()
        d.goresym = "/tmp/goresym.json"
        with patch("builtins.open", mock_open(read_data="not json {")):
            d._BinaryNinjaDecompiler__apply_goresym()
            # Should log error but not crash
            d.log.log_error if hasattr(d.log, 'log_error') else None
            # Just verify no exception raised