import datetime
import json
import os
import sys
import time

from typing import Dict, Optional, List, Tuple, Literal, Any, Union, Annotated, TypedDict

import malcat
malcat.setup()
from zeromcp.mcp import McpServer, McpToolError
from transform import get_transforms


class RegionInfo(TypedDict):
    name: Annotated[str, "region name"]
    effective_address: Annotated[int, "effective address of first byte of the region"]
    physical_size: Annotated[int, "size in bytes of the region on disk"]
    virtual_size: Annotated[int, "size in bytes of the region in memory"]
    rights: Annotated[str, "rights of region"]
    entropy: Annotated[int, "entropy of region, from 0 (low) to 255 (high)"]

class KesakodeVerdict(TypedDict):
    malware_family: Annotated[str, "name of malware family"]
    confidence: Annotated[str, "signature confidence"]

class MalcatAnalysisEntry(TypedDict):
    analysis_id: Annotated[int, "unique identifier for this analysis"]
    file_path: Annotated[str, "path of analysed file on disk, if relevant"]
    file_name: Annotated[str, "name of analysed file"]
    type: Annotated[str, "identified file type"]

class MalcatAnalysisResult(MalcatAnalysisEntry):
    file_size: Annotated[int, "file size"]
    architecture: Annotated[str, "CPU architecture (if program)"]
    entropy: Annotated[int, "entropy of file, from low (0) to high (255)"]
    sha256: Annotated[str, "sha256 of file"]
    metadata: Annotated[Dict[str, str], "file's extracted metadata"]
    entrypoint_ea: Annotated[int, "effective address of entry point (if program)"]
    layout: Annotated[List[RegionInfo], "layout of the sections/regions"]
    kesakode_verdict: Annotated[List[KesakodeVerdict], "Kesakode's malware verdict. Kesakode can have false positives"]

    @staticmethod
    def from_analysis(analysis: malcat.Analysis):

        def get_layout(a) -> List[RegionInfo]:
            res = []
            for region in a.map:
                if region.phys_size:
                    entropy = a.entropy[region.phys: region.phys + region.phys_size]
                else:
                    entropy = 0
                res.append(RegionInfo(
                    name=region.name,
                    effective_address=region.address,
                    physical_size=region.phys_size,
                    virtual_size=region.virt_size, 
                    rights=(region.read and "R" or "") + (region.write and "W" or "") + (region.exec and "X" or ""),
                    entropy=entropy
                ))
            return res


        def get_kesakode_verdict(analysis) -> List[KesakodeVerdict]:
            res = []
            for malware_family, score in sorted(analysis.kesakode.verdict.items(), key=lambda x: -x[1])[:4]:
                res.append(KesakodeVerdict(malware_family=malware_family, confidence=f"{int(score):d}/100"))
            return res

        metadata = {}
        for cat, ms in analysis.metadata:
            for k, v in ms:
                if cat:
                    k = f"{cat}::{k}"
                    metadata[k] = v
        return MalcatAnalysisResult(
            analysis_id=analysis.id, 
            file_name=analysis.file.name, 
            file_path=analysis.file.path, 
            file_size=analysis.file.size,
            type=analysis.type,
            architecture=analysis.architecture.name,
            entropy=analysis.entropy[0:len(analysis.file)],
            sha256=analysis.entropy.sha256,
            metadata=metadata,
            entrypoint_ea=analysis.entrypoint,
            layout=get_layout(analysis),
            kesakode_verdict=get_kesakode_verdict(analysis)
            )



class MalcatMCPServer(McpServer): 
    """Manages Malcat binary analyses"""

    def __init__(self, max_binaries: int = 10, kesakode_key: str = ""):
        McpServer.__init__(self, "malcat MCP server", version="1.0.0")
        self._analyses: Dict[int, malcat.Analysis] = {}
        self._load_times: Dict[str, float] = {}
        self.max_binaries = max_binaries
        self.kesakode_key = kesakode_key

    def register_analysis(self, analysis: malcat.Analysis):
        if len(self._analyses) >= self.max_binaries:
            oldest = min(self._load_times, key=self._load_times.get)
            del self._analyses[oldest]
            del self._load_times[oldest]
        self._analyses[analysis.id] = analysis
        self._load_times[analysis.id] = time.time()


    def analyse_file(self, malcat_file: malcat.File) -> int:
        """Load and analyse a binary"""
        analysis = malcat.analyse(malcat_file)
        analysis.raise_if_failed()
        if len(analysis.fns) and self.kesakode_key:
            analysis.kesakode.online_lookup(self.kesakode_key)
        self.register_analysis(analysis)
        return analysis

    def get_analysis(self, analysis_id:Annotated[int, "analysis ID"]) -> malcat.Analysis:
        """Get a loaded analysis"""
        if analysis_id not in self._analyses:
            raise McpToolError(f"Analysis #{analysis_id} not found, did you analyse the file first with analyse_file?")
        return self._analyses[analysis_id]

    def list_analyses(self) -> List[MalcatAnalysisEntry]:
        """List all analysises in context"""
        res = []
        for a in self._analyses.values():
            res.append(MalcatAnalysisEntry(id=a.id, file_name=a.file.name, file_path=a.file.path, type=a.type))
        return res

    def delete_analysis(self, analysis_id:Annotated[int, "analysis ID"]) -> bool:
        """Unload a binary"""
        if analysis_id in self._analyses:
            del self._analyses[analysis_id]
            if analysis_id in self._load_times:
                del self._load_times[analysis_id]
            return True
        return False

    def close(self):
        """Cleanup"""
        self._analyses.clear()
        self._load_times.clear()


def smart_number_convert(value):
    if type(value) == str:
        if value.startswith("0x"):
            value = int(value, 16)
        elif value.startswith("#"):
            value = int(value[1:], 16)
        else:
            value = int(value)
    return value

def smart_ea_convert(analysis, value):
    original_value = value
    if type(value) == str:
        if value.startswith("0x"):
            value = int(value, 16)
            # chances are model messed up and it is a va
            from_va = analysis.v2a(value)
            if from_va is not None:
                value = from_va
        elif value.startswith("#"):
            # chances are model messed up and it is an offset
            value = int(value[1:], 16)
            from_pa = analysis.p2a(value)
            if from_pa is not None:
                value = from_pa
        elif value.startswith("@"):
            # chances are model messed up and it is a rva
            value = int(value[1:], 16)
            from_rva = analysis.r2a(value)
            if from_rva is not None:
                value = from_rva
        else:
            value = int(value)
    if value not in analysis.map:
        raise McpToolError(f"address {original_value} is not a valid effective address")
    return value

# ============================================================================
# Tools
# ============================================================================

def register_tools_common(mcp: MalcatMCPServer):
    """Register all tools"""


    

    ################################# ANALYSIS

    @mcp.tool
    def analyse_forget(analysis_id:Annotated[int, "analysis ID"]):
        """
        Discard a given analysis
        """
        if mcp.delete_analysis(analysis_id):
            return "OK"
        else:
            return "Failed"

    @mcp.tool
    def analyse_infos(analysis_id:Annotated[int, "analysis ID"]) -> MalcatAnalysisResult:
        """
        Display basic infos about an analysis
        """
        analysis = mcp.get_analysis(analysis_id)
        return MalcatAnalysisResult.from_analysis(analysis)

    @mcp.tool
    def analyse_file(path:Annotated[str, "path of file on disk"]) -> MalcatAnalysisResult:
        """
        load a file from disk and analyse it for further inspection/decryption. Returns a short summary of the analysis.
        """
        path = os.path.abspath(os.path.expanduser(os.path.normpath(path)))
        if not os.path.exists(path):
            raise FileNotFoundError(f"File not found: {path}")
        mf = malcat.File.open(path)
        return MalcatAnalysisResult.from_analysis(mcp.analyse_file(mf))


    @mcp.tool
    def analyse_file_interval(analysis_id:Annotated[int, "analysis ID"], ea:Annotated[Union[int,str], "effective address of first byte to analyse"], size:Annotated[int, "how many bytes to read"]) -> MalcatAnalysisResult:
        """
        load file interval and analyse it for further inspection/decryption. Returns a short summary of the analysis.
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        offset = analysis.a2p(ea)
        if offset is None:
            raise McpToolError("effective address is not file-backed (cannot convcert to offset)")
        mf = malcat.File.subfile(analysis.file, offset, size, f"analysis_{analysis_id}_interval_{ea}_{ea+size}")
        return MalcatAnalysisResult.from_analysis(mcp.analyse_file(mf))

    @mcp.tool
    def analyse_carved_file(analysis_id:Annotated[int, "analysis ID where the carved file was found"], ea:Annotated[Union[int,str], "effective addresse where the file was carved"]) -> MalcatAnalysisResult:
        """
        load a carved file previously found at <ea> and analyse it for further inspection/decryption. Returns a short summary of the analysis.
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        carved_file = analysis.carved[ea]
        return MalcatAnalysisResult.from_analysis(mcp.analyse_file(carved_file.open()))

    @mcp.tool
    def analyse_virtual_file(analysis_id:Annotated[int, "analysis ID containing the virtual file"], vfile_path:Annotated[str, "virtual file path"], password:Annotated[str, "file password, leave empty for best guessed pwd"]="") -> MalcatAnalysisResult:
        """
        load a virtual file of path <vfile_path> and analyse it for further inspection/decryption/decryption. Returns a short summary of the analysis.
        """
        analysis = mcp.get_analysis(analysis_id)
        f = analysis.open_vfile(vfile_path, password)
        return MalcatAnalysisResult.from_analysis(mcp.analyse_file(f))

    @mcp.tool
    def pretty_print_ea(analysis_id:Annotated[int, "analysis ID"], ea:Annotated[Union[int,str], "effective address"]) -> str:
        """
        Pretty-print an effective address
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        return analysis.ppa(ea)


    @mcp.tool
    def ea_to_offset(analysis_id:Annotated[int, "analysis ID"], ea:Annotated[Union[int,str], "effective address"]) -> Optional[int]:
        """
        Convert an effective address <ea> to a file offset
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        r = analysis.a2p(ea)
        if r is None:
            raise McpToolError(f"{ea} is not a valid effective address, maybe a virtual address?")
        return r

    @mcp.tool
    def ea_to_va(analysis_id:Annotated[int, "analysis ID"], ea:Annotated[Union[int,str], "effective address"]) -> Optional[int]:
        """
        Convert an effective address to a virtual address
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        r = analysis.a2v(ea)     
        if r is None:
            raise McpToolError(f"{ea} is not a valid effective address, maybe a virtual address?")
        return r

    @mcp.tool
    def offset_to_ea(analysis_id:Annotated[int, "analysis ID"], offset:Annotated[Union[int,str], "offset"]) -> Optional[int]:
        """
        Convert a file offset to an effective address
        """
        analysis = mcp.get_analysis(analysis_id)
        offset = smart_number_convert(offset)
        r = analysis.p2a(offset)
        if r is None:
            raise McpToolError(f"{offset} is not a valid offset")
        return r

    @mcp.tool
    def va_to_ea(analysis_id:Annotated[int, "analysis ID"], va:Annotated[Union[int,str], "virtual address"]) -> Optional[int]:
        """
        Convert a virtual address to an effective address
        """
        analysis = mcp.get_analysis(analysis_id)
        va = smart_number_convert(va)
        r = analysis.v2a(va)               
        if r is None:
            raise McpToolError(f"{va} is not a valid virtual address")
        return r

    @mcp.tool
    def va_to_offset(analysis_id:Annotated[int, "analysis ID"], va:Annotated[Union[int,str], "virtual address"]) -> Optional[int]:
        """
        Convert a virtual address to an offset in file
        """
        analysis = mcp.get_analysis(analysis_id)
        va = smart_number_convert(va)
        r = analysis.v2a(va)       
        if r is None:
            raise McpToolError(f"{va} is not a valid virtual address")
        r = analysis.a2p(r)
        if r is None:
            raise McpToolError(f"{va} is not file-backed")
        return r

    ################################# FILE

    class CarvedFile(TypedDict):
        ea: Annotated[int, "effective address where the file was found"]
        offset: Annotated[int, "offset on disk where the file was found"]
        size: Annotated[int, "size of file in bytes"]
        type: Annotated[str, "type of carved file"]
        type_category: Annotated[str, "category of carved file"]

    @mcp.tool
    def file_list_carved(analysis_id:Annotated[int, "analysis ID"]) -> List[CarvedFile]:
        """
        list all files of know type found via file carving
        """
        analysis = mcp.get_analysis(analysis_id)
        res = []
        for f in analysis.carved:
            res.append(CarvedFile(
                ea=f.address,
                offset=analysis.a2p(f.address),
                size=f.size,
                type=f.type,
                type_category=f.category.name,
            ))
        return res


    class VirtualFile(TypedDict):
        path: Annotated[str, "virtual file path"]
        unpacked_size: Annotated[int, "size of virtual file in bytes"]

    @mcp.tool
    def file_list_virtual_files(analysis_id:Annotated[int, "analysis ID"]) -> List[VirtualFile]:
        """
        list all virtual files (archive members, program resources, etc.) stored inside the file
        """
        analysis = mcp.get_analysis(analysis_id)
        res = []
        for f in analysis.vfiles:
            res.append(VirtualFile(path=f.path, unpacked_size=f.size))
        return res

    @mcp.tool
    def save_virtual_file_to_disk(analysis_id:Annotated[int, "analysis ID"], vfile_path: Annotated[str, "path of virtual file"], output_path:Annotated[str, "where to write the result on disk"]):
        """
        Save a virtual file to disk
        """
        analysis = mcp.get_analysis(analysis_id)
        f = analysis.open_vfile(vfile_path)
        dirname = os.path.dirname(output_path)
        if not os.path.exists(dirname):
            os.makedirs(dirname)        
        if not f.save(output_path):
            raise McpToolError("File could not be saved to disk")

    @mcp.tool
    def save_file_interval_to_disk(analysis_id:Annotated[int, "analysis ID"], ea:Annotated[Union[int, str], "effective address of the first byte of the file to dump"], size:Annotated[int, "how many bytes to dump. 0 = until end of file"], output_path:Annotated[str, "where to write the result on disk"]):
        """
        Save a file interval to disk
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        offset = analysis.a2p(ea)
        if offset is None:
            raise McpToolError("effective address is not file-backed (cannot convcert to offset)")
        mf = malcat.File.subfile(analysis.file, offset, size, f"analysis_{analysis_id}_interval_{offset}_{offset+size}")
        dirname = os.path.dirname(output_path)
        if not os.path.exists(dirname):
            os.makedirs(dirname)
        if not mf.save(output_path):
            raise McpToolError("File interval could not be saved to disk")

    @mcp.tool
    def file_read(analysis_id:Annotated[int, "analysis ID"], ea:Annotated[Union[int, str], "effective address of the first byte of the file to read"], size:Annotated[int, "how many bytes to read"] = 1024) -> str:
        """
        Read raw bytes in the file and return them hex-encoded
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        offset = analysis.a2p(ea)
        if offset is None:
            raise McpToolError("effective address is not file-backed (cannot convcert to offset)")
        offset = smart_number_convert(offset)
        if size > 1024:
            raise McpToolError("Too much data to read, would fill context")
        return analysis.file[offset:offset + size].hex().upper()


    @mcp.tool
    def file_interval_entropy(analysis_id:Annotated[int, "analysis ID"], ea:Annotated[Union[int, str], "effective address of the first byte of the file interval"], size:Annotated[int, "how many bytes"]) -> int:
        """
        Compute the entropy of a given file interval as a number between 0 (low) and 255 (very high)
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        offset = analysis.a2p(ea)
        if offset is None:
            raise McpToolError("effective address is not file-backed (cannot convcert to offset)")
        return analysis.entropy[offset:offset + size]


    @mcp.tool
    def script_decompile(analysis_id:Annotated[int, "analysis ID"], output_path:Annotated[str, "where to write the result on disk"]):
        """
        decompile a script and write it to disk. Supported file types: 
        - AU3.Tokens -> recovers textual AutoIT script
        - CFB -> recovers MSI tables if MSI installer or VBA code if PPT/WORD/EXCEL
        - Office.Workbook8/12 -> recovers Excel macros
        - VBE: recovers VBS script
        - PascalScript/NSIS.Setup: disassemble all functions
        """
        analysis = mcp.get_analysis(analysis_id)
        
        dirname = os.path.dirname(output_path)
        if not os.path.exists(dirname):
            os.makedirs(dirname)        
        if analysis.type == "VBE":
            data = analysis.open_vfile("<content>")[:].decode("utf8", errors="replace")
        elif analysis.architecture in (malcat.Architecture.NSIS, malcat.Architecture.PASCALSCRIPT):
            data = "\n\n".join([fn.disasm() for fn in analysis.fns])
        elif analysis.type in ("AU3.Tokens", "CFB", "Office.Workbook12", "Office.Workbook8"):
            data = analysis.decompile()
        else:
            raise McpToolError(f"unsupported file type {analysis.type}")
        with open(output_path, "w", encoding="utf8") as f:
            f.write(data)

    ################################# FUNCTIONS

    def function_score(fn):
        return (
                fn.category.value in (malcat.Detection.Level.MALWARE, malcat.Detection.Level.UNKNOWN) and 1 or 0,
                fn.num_known_constants_uses, 
                fn.num_xor_ops, 
                fn.num_api_calls,
                fn.num_string_references,
                fn.num_unique_immediate_bytes,
                fn.num_arithmetic_ops,
                )


    class FunctionHit(TypedDict):
        ea: Annotated[int, "effective address of function"]
        name: Annotated[str, "sfunction name"]        

    class StringHit(TypedDict):
        ea: Annotated[int, "effective address of string"]
        summary: Annotated[str, "the string's content, truncated to 32 chars"]        


    @mcp.tool
    def fns_top_list(analysis_id:Annotated[int, "analysis ID"], maximum_number_of_functions:Annotated[int, "max to return"]=128) -> List[FunctionHit]:
        """
        list the most interesting functions.
        """
        analysis = mcp.get_analysis(analysis_id)
        all_fns = sorted(list(analysis.fns), key=lambda fn: function_score(fn), reverse=True)
        res = []
        for fn in all_fns[:maximum_number_of_functions or len(all_fns)]:
            res.append(FunctionHit(ea=fn.address, name=fn.name))
        return res

    @mcp.tool
    def fns_search(analysis_id:Annotated[int, "analysis ID"], maximum_number_of_functions:Annotated[int, "max to return"]=128,
            contains:Annotated[str, "substring which must be found anywhere inside the function prototype"]="", 
            starts_with:Annotated[str, "substring which must be found at the beginning of the function prototype"]="", 
            ends_with:Annotated[str, "substring which must be found at the end of the function prototype"]="") -> List[FunctionHit]:
        """
        search among all identified strings following _all_ the given search criterions.
        """
        analysis = mcp.get_analysis(analysis_id)
        matching = []
        for fn in analysis.fns:
            if contains and contains not in fn.fullname:
                continue
            if starts_with and not (fn.name.startswith(starts_with) or fn.fullname.startswith(starts_with)):
                continue
            if ends_with and not (fn.name.endswith(ends_with) or fn.fullname.endswith(ends_with)):
                continue
            matching.append(fn)
        res = []
        for fn in sorted(matching, key=lambda fn: function_score(fn), reverse=True)[:maximum_number_of_functions]:
            res.append(FunctionHit(ea=fn.address, name=fn.name))
        return res

    class FunctionStats(TypedDict):
        name:Annotated[str, "function's name (without module name]"]
        prototype:Annotated[str, "full function prototype"]
        start:Annotated[int, "effective address of the function's first byte"]
        size:Annotated[int, "function's size in bytes"]
        category:Annotated[str, "CLEAN, LIBRARY, MALWARE (only seen in malware) or UNKNOWN"]
        num_basic_blocks:Annotated[int, "number of basic blocks"]
        num_instructions:Annotated[int, "number of instructions"]
        num_intra_jumps:Annotated[int, "number of jumps pointing inside the function"]
        num_dyn_jumps:Annotated[int, "number of jump <register>"]
        num_api_calls:Annotated[int, "number of direct API calls"]
        num_string_refs:Annotated[int, "number of string referenced in the function"]
        num_known_constants:Annotated[int, "number of know constant uses seen"]
        num_arithmetic_ops:Annotated[int, "number of add/sub/div/mul-like opcodes"]
        num_assign_ops:Annotated[int, "number of mov-like opcodes"]
        num_binary_ops:Annotated[int, "number of and/or/xor/shl/shr-like opcodes"]
        num_xor_ops:Annotated[int, "number of xor-like opcodes"]
        referenced_strings:Annotated[List[StringHit], "list of referenced strings"]

    @mcp.tool
    def fn_infos(analysis_id:Annotated[int, "analysis ID"], ea:Annotated[Union[int,str], "effective address of function"]) -> FunctionStats:
        """
        return some information about a function, such as the number of basic blocks, referenced strings, how many api calls, etc.
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        fn = analysis.fns[ea]
        strs = []
        for o in fn.outrefs:
            s = analysis.strings.find(o.address)
            if s:
                strs.append(StringHit(ea=s.address, summary=s.summary))
        return FunctionStats(
                name = fn.name,
                prototype = fn.fullname,
                start = fn.start,
                size = len(fn),
                category = fn.category.name,
                num_basic_blocks = fn.num_bb_code,
                num_instructions = fn.num_instructions,
                num_intra_jumps = fn.num_intra_jumps,
                num_dyn_jumps = fn.num_dynamic_jumps,
                num_highvalue_immediates = fn.num_highvalue_immediates,
                num_api_calls = fn.num_api_calls,
                num_string_refs = fn.num_string_references,
                num_known_constants = fn.num_known_constants_uses,
                num_arithmetic_ops = fn.num_arithmetic_ops,
                num_assign_ops = fn.num_assign_ops,
                num_binary_ops = fn.num_binary_ops,
                num_xor_ops = fn.num_xor_ops,
                referenced_strings = strs
                )


    class FunctionCallersCallees(TypedDict):
        callers: Annotated[List[FunctionHit], "list of callers"]
        callees: Annotated[str, "list of callees"]        
    @mcp.tool
    def fn_callers_callees(analysis_id:Annotated[int, "analysis ID"], ea:Annotated[Union[int,str], "effective address of function"]) -> List[FunctionHit]:
        """
        list the functions called by/calling to a given function
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        fn = analysis.fns[ea]
        return FunctionCallersCallees(
            callers=[FunctionHit(ea=f.address, name=f.name) for f in fn.callers],
            callees=[FunctionHit(ea=f.address, name=f.name) for f in fn.callees]
        )
            
    @mcp.tool
    def fn_disassemble(analysis_id:Annotated[int, "analysis ID"], ea:Annotated[Union[int,str], "effective address of function to inspect"]) -> str:
        """
        disassemble a given function, for architectures: x86/x64, arm, aarch64, mips, PYC, .NET, PascalScript and NSIS. 
        Note that every address in the disassembly listing are virtual addresses.
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        fn = analysis.fns[ea]
        return fn.disasm()


    @mcp.tool
    def fn_decompile(analysis_id:Annotated[int, "analysis ID"], ea:Annotated[Union[int,str], "effective address of function to inspect"]) -> str:
        """
        decompile a given function. Only works for x86, x64, mips, arm and aarch64 architectures. If decompiling does not work, you can try to disassemble the function instead.
        Note that every address that you see there are virtual addresses.
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        fn = analysis.fns[ea]
        return fn.decompile()


    ################################# STRINGS


    @mcp.tool
    def strings_top_list(analysis_id:Annotated[int, "analysis ID"], maximum_number_of_strings:Annotated[int, "maximum number of strings to return"]=128) -> List[StringHit]:
        """
        list the most interesting strings found in file
        """
        analysis = mcp.get_analysis(analysis_id)
        maximum_number_of_strings = min(maximum_number_of_strings, 1000)
        res = []
        for i, s in enumerate(analysis.strings.filter()):
            if i >= maximum_number_of_strings:
                break
            res.append(StringHit(ea=s.address, summary=s.summary))
        return res

    @mcp.tool
    def strings_search(analysis_id:Annotated[int, "analysis ID"], 
            maximum_number_of_strings:Annotated[int, "maximum number of strings to return"]=128,
            contains:Annotated[str, "substring which must be found anywhere inside the string"]="", 
            starts_with:Annotated[str, "substring which must be found at the beginning of the string"]="", 
            ends_with:Annotated[str, "substring which must be found at the end of the stringe"]="",
            type:Annotated[str, '"USER" for user-defined strings, "META" for compiler-defined, "DYNAMIC" for stack strings or "" for all']="",
            tag:Annotated[str, 'tag that the string must have. One of "BASE64", "HEXA, "IP", "PATH", "URL" or ""']="",
            size_min:Annotated[int, "minimum size of the string in characters"] = 0,
            score_min:Annotated[int, " minimum score of the string, between 0 (uninteresting) and 255 (most interesting)"] = 80,
         ) -> List[StringHit]:
        """
        search among all identified strings following the given criterions.
        """
        analysis = mcp.get_analysis(analysis_id)
        maximum_number_of_strings = min(maximum_number_of_strings, 1000)
        res = []
        if type:
            type = {
                "USER": malcat.FoundString.Type.USER,
                "META": malcat.FoundString.Type.META,
                "DYNAMIC": malcat.FoundString.Type.DYNAMIC,
                "SCANNED": malcat.FoundString.Type.SCANNED,
            }.get(type, None)
        else:
            type = None
        for i, s in enumerate(analysis.strings.filter(
                contains=contains,
                starts_with=starts_with,
                ends_with=ends_with,
                type=type,
                tag=tag,
                size_min=size_min,
                score_min=score_min,
                )):
            if i >= maximum_number_of_strings:
                break
            res.append(StringHit(ea=s.address, summary=s.summary))
        return res


    class StringStats(TypedDict):
        text: Annotated[str, "first 512 chars of the string"]
        type: Annotated[str, "USER (user-defined string), META (compiler-defined string), DYNAMIC (stack-string) or SCANNED (unreferenced string)"]
        encoding: Annotated[str, "string encoding"]
        tag: Annotated[str, "string tag"]
        num_characters: Annotated[int, "number of characters"]
        size: Annotated[int, "size in bytes"]
        score: Annotated[int, "how much ist the string interesting (0-255)"]

    @mcp.tool
    def string_infos(analysis_id:Annotated[int, "analysis ID"], ea:Annotated[Union[int,str], "effective address of string"]) -> StringStats:
        """
        return some information about a string, such as its full text, encoding, type or score.
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        s = analysis.strings[ea]
        return StringStats(
                text = s.text[:512],
                type = s.type.name,
                encoding = s.encoding.name,
                tag = s.tag or "",
                num_characters = s.characters,
                size = s.size,
                score = s.score,
            )

    ################################# ANOMALIES

    class AnomalyInfo(TypedDict):
        name: Annotated[str, "anomaly internal name"]
        desc: Annotated[str, "description"]
        category: Annotated[str, "anomaly category"]
        level: Annotated[int, "how odd is the anomaly found, from 1 (slightly odd) to 4 (very odd)"]
        num_hits: Annotated[int, "number of times it was found in the file"]

    @mcp.tool
    def anomalies_list(analysis_id:Annotated[int, "analysis ID"]) -> List[AnomalyInfo]:
        """
        list all the anomalies and oddities worth of interest found inside the file
        """
        analysis = mcp.get_analysis(analysis_id)
        res = []
        for a in analysis.anomalies:
            res.append(AnomalyInfo(
                name = a.name,
                desc = a.comment,
                category = a.category,
                level = a.level.value + 1,
                num_hits = len(a.locations),
            ))
        return res

    class AnomalyHit(TypedDict):
        ea: Annotated[int, "anomaly effective address"]
        size: Annotated[int, "anomaly size"]

    @mcp.tool
    def anomaly_list_locations(analysis_id:Annotated[int, "analysis ID"], anomaly_name: Annotated[str, "anomaly internal name"]) -> List[AnomalyHit]:
        """
        list all occurences where the anomaly has been found in the file
        """
        analysis = mcp.get_analysis(analysis_id)
        return [AnomalyHit(ea=x[0], size=x[1]) for x in analysis.anomalies[anomaly_name].locations]

    ################################# XREFS

    class IncomingReference(TypedDict):
        source_ea: Annotated[int, "source effective address"]
        type: Annotated[str, 'either "code" for code references or "data" for data pointers in a data section']
        source_function_ea: Annotated[Optional[int], "if code reference, effective address of the function containing the referencing instruction"]
        source_function_name: Annotated[Optional[str], "if code reference, name of the function containing the referencing instruction"]

    @mcp.tool
    def refs_get(analysis_id:Annotated[int, "analysis ID"], ea:Annotated[Union[int,str], "refs target (effective address)"]) -> List[IncomingReference]:
        """
        list all incoming references to <ea>
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        res = []
        for r in analysis.xrefs[ea]:
            d = IncomingReference(source_ea=r.address, type = r.is_code and "code" or "data")
            fn = analysis.fns.find(r.address)
            if fn is not None:
                d["source_function_ea"] = fn.address
                d["source_function_name"] = fn.name
            res.append(d)
        return res

    ################################# MATCH

    class ConstantInfo(TypedDict):
        id: Annotated[str, "constant id"]
        category: Annotated[str, "category"]        
        type: Annotated[str, "'code' for immediates found in code, or 'data'"]        
        num_hits: Annotated[int, "number of times it was found in the file"]        

    def get_all_constants_by_id(analysis:malcat.Analysis):
        grouped_by_id = {}
        for c in analysis.constants:
            id = f"{c.category}::{c.name}"
            grouped_by_id.setdefault(id, []).append(c)
        return grouped_by_id


    @mcp.tool
    def constants_list(analysis_id:Annotated[int, "analysis ID"]) -> List[ConstantInfo]:
        """
        list all the know constants found inside data and/or code
        """
        analysis = mcp.get_analysis(analysis_id)
        res = []
        for k, v in get_all_constants_by_id(analysis).items():
            res.append( ConstantInfo(id=k, category=v[0].category, type=v[0].is_code and "code" or "data", num_hits=len(v)))
        return res


    class ConstantHit(TypedDict):
        ea: Annotated[int, "constant effective address"]
        size: Annotated[int, "constant size"]

    @mcp.tool
    def constant_list_locations(analysis_id:Annotated[int, "analysis ID"], constant_id: Annotated[str, "constant id"]) -> List[ConstantHit]:
        """
        list all occurences where the constant has been found in the file
        """
        analysis = mcp.get_analysis(analysis_id)
        grouped = get_all_constants_by_id(analysis)
        res = []
        for c in grouped.get(constant_id, []):
            res.append(ConstantHit(ea=c.address, size=c.size))
        return res


    class YaraMatch(TypedDict):
        id: Annotated[str, "Yara rule id"]
        category: Annotated[str, "Yara rule category"]
        reliability: Annotated[int, "0 (very FP-prone) to 100 (no FP at all)"]
        type: Annotated[str, "dangerosity of the rule"]
        description: Annotated[str, "what the rule detects"]
        num_patterns: Annotated[int, "number of matching string/byte patterns"]

    @mcp.tool
    def yara_list(analysis_id:Annotated[int, "analysis ID"]) -> List[YaraMatch]:
        """
        list all matching Yara signatures
        """
        analysis = mcp.get_analysis(analysis_id)
        res = []
        for m in analysis.sigs:
            res.append( YaraMatch(
                id = m.id, 
                category = m.category, 
                reliability = m.reliability, 
                type = m.type.name, 
                description = m.description, 
                num_patterns = len(m)
                ))
        return res


    class YaraPatternHit(TypedDict):
        name: Annotated[str, "pattern name"]
        ea: Annotated[int, "pattern hit effective address"]
        size: Annotated[int, "pattern hit size"]

    @mcp.tool
    def yara_list_locations(analysis_id:Annotated[int, "analysis ID"], yara_rule_id: str) -> List[YaraPatternHit]:
        """
        list the locations of all the matching string/byte patterns for the Yara rule
        """
        analysis = mcp.get_analysis(analysis_id)
        res = []
        for p in analysis.sigs[yara_rule_id].patterns:
            for off, size in p.matches:
                res.append(YaraPatternHit(name=p.id, ea=analysis.p2a(off), size=size))
        return res    

    ################################# TRANSFORMS

    ALL_TRANSFORMS = { trans.name: (trans, trans.get_infos()) for trans in get_transforms(malcat.env.datadir, malcat.env.userdir) }


    class TransformParameter(TypedDict):
        type: Annotated[str, "parameter type: BOOL, INT, STRING, ENUM or DATA. Data means hex-encoded bytes."]
        possible_values: Annotated[Optional[List[str]], "list of possible values, for the ENUM type"] = None
        default_value: Annotated[Optional[Any], "default value"] = None

    class TransformInfo(TypedDict):
        name: Annotated[str, "Transform name"]
        category: Annotated[str, "Category"]
        description: Annotated[str, "Description"]
        parameters: Annotated[Dict[str, TransformParameter], "transform prototype"]

    def list_all_transforms() -> List[TransformInfo]:
        """
        list all transforms supported by Malcat
        """
        res = []
        for transf, infos in ALL_TRANSFORMS.values():
            if "script" in infos.category or "python" in infos.name:
                continue
            res.append(TransformInfo(
                name=infos.name,
                category=infos.category,
                description=infos.doc,
                parameters={ x.name: TransformParameter(type=x.type.name, possible_values=x.possible_values, default_value=x.default_value) for x in infos.parameters }
            ))
            res[-1]["parameters"]["input_data"] = TransformParameter(type="DATA")
        return res


    def convert_transform_params(transform_infos, transform_parameters):
        params = {}
        
        for p in transform_infos.parameters:
            if not p.name in transform_parameters:
                pass#raise ValueError(f"Parameter {p.name} was not specified")
            else:
                input_val = transform_parameters[p.name]
                if p.type == malcat.TransformParameter.Type.BOOL:
                    if type(input_val) != bool:
                        raise McpToolError(f"Parameter {p.name} should be of type bool but was {type(input_val)}")
                elif p.type == malcat.TransformParameter.Type.INT:
                    if type(input_val) != int:
                        raise McpToolError(f"Parameter {p.name} should be of type int but was {type(input_val)}")
                elif p.type == malcat.TransformParameter.Type.STRING:
                    if type(input_val) != str:
                        raise McpToolError(f"Parameter {p.name} should be of type str but was {type(input_val)}")
                elif p.type == malcat.TransformParameter.Type.ENUM:
                    if type(input_val) != str:
                        raise McpToolError(f"Parameter {p.name} should be of type str, and one of {p.possible_values}")
                    elif input_val not in p.possible_values:
                        raise McpToolError(f"Parameter {p.name} should be one of {p.possible_values}")
                elif p.type == malcat.TransformParameter.Type.DATA:
                    try:
                        input_val = bytes.fromhex(input_val)
                    except BaseException as e:
                        raise McpToolError(f"Parameter {p.name} should be a hex-encoded byte string: {e}")
                params[p.name] = input_val
                del transform_parameters[p.name]

        for k in transform_parameters:
            raise McpToolError(f"Unknown parameter {k}")

        return params


    @mcp.tool
    def transforms_search(algo_name:Annotated[str, "(partial) name of algorithm to search"]) -> List[TransformInfo]:
        """
        search the list of supported decryption/decoding algorithm/transforms and return their name and prototype
        """
        to_search = [x.lower() for x in algo_name.replace("_", " ").split()]
        res = []
        for tinfo in list_all_transforms():
            for q in to_search:
                if q in tinfo["name"]:
                    res.append(tinfo)
                    break
        return res




    @mcp.tool
    def hex_encode_string(input_string:Annotated[str, "the string to hex-encode"], encoding:Annotated[str, "encoding of the string: 'utf-8' or 'utf-16le'"]="utf-8") -> str:
        """
        Encode a string using hexadecimal-encoding.
        """
        return input_string.encode(encoding).hex().upper()
    
    class TransformPass(TypedDict):
        name: Annotated[str, "the name of the transform to apply (cf. transforms_list)"]
        parameters: Annotated[Dict[str, Any], "additional parameters for the transform. Parameters of type DATA must be given as hex-encoded string"]

    @mcp.tool
    def decrypt_hexencoded_buffer(
            input_data:Annotated[str, "hex-encoded input data"], 
            list_of_transforms:Annotated[List[TransformPass], "the sequence of transforms to apply"] 
        ) -> MalcatAnalysisResult:
        """
        Decrypt the given (short) hex-encoded input data by applying the given <list_of_transforms> sequentially
        result is loaded in a new analysis
        """
        try:
            data = bytes.fromhex(input_data)
        except BaseException as e:
            raise McpToolError(f"input_data should be a hex-encoded byte string: {e}")
        for t in list_of_transforms:
            transform_name = t["name"]
            try:
                trans, infos = ALL_TRANSFORMS[transform_name]
                params = convert_transform_params(infos, t["parameters"])
                data = trans().run(data, **params)
            except BaseException as e:
                raise McpToolError(f"Error in transfor {transform_name}: {e}")
        mf = malcat.File.from_buffer(data)
        return MalcatAnalysisResult.from_analysis(mcp.analyse_file(mf))

    @mcp.tool
    def decrypt_string(
            analysis_id:Annotated[int, "analysis ID"], 
            ea:Annotated[Union[int,str], "effective addresse of string"],
            list_of_transforms:Annotated[List[TransformPass], "the sequence of transforms to apply"] 
        ) -> MalcatAnalysisResult:
        """
        Decrypt/decode the given string by applying the given <list_of_transforms> sequentially
        result is loaded in a new analysis
        """
        analysis = mcp.get_analysis(analysis_id)
        ea = smart_ea_convert(analysis, ea)
        s = analysis.strings.find(ea)
        if not s:
            raise McpToolError("Could not find string at address")
        data = s.text.encode("utf-8")
        for t in list_of_transforms:
            transform_name = t["name"]
            try:
                trans, infos = ALL_TRANSFORMS[transform_name]
                params = convert_transform_params(infos, t["parameters"])
                data = trans().run(data, **params)
            except BaseException as e:
                raise McpToolError(f"Error in transfor {transform_name}: {e}")
        mf = malcat.File.from_buffer(data)
        return MalcatAnalysisResult.from_analysis(mcp.analyse_file(mf))

    @mcp.tool
    def chain_decrypt_analysis(
            analysis_id:Annotated[int, "analysis ID of the file to decrypt"],
            list_of_transforms:Annotated[List[TransformPass], "the sequence of transforms to apply"] 
        ) -> MalcatAnalysisResult:
        """
        Decode/decrypt a whole file by applying the given <list_of_transforms> sequentially
        result is loaded in a new analysis
        """
        analysis = mcp.get_analysis(analysis_id)
        data = analysis.file[:]
        
        for t in list_of_transforms:
            transform_name = t["name"]
            try:
                trans, infos = ALL_TRANSFORMS[transform_name]
                params = convert_transform_params(infos, t["parameters"])
                data = trans().run(data, **params)
            except BaseException as e:
                raise McpToolError(f"Error in transfor {transform_name}: {e}")
        mf = malcat.File.from_buffer(data)
        return MalcatAnalysisResult.from_analysis(mcp.analyse_file(mf))

    ################################# STRUCTURES

    class StructInfo(TypedDict):
        name: Annotated[str, ""]
        ea: Annotated[str, "effective address"]        

    @mcp.tool
    def structs_list(analysis_id:Annotated[int, "analysis ID"]) -> List[StructInfo]:
        """
        list all parsed file structures
        """
        analysis = mcp.get_analysis(analysis_id)
        return [StructInfo(name=s.name, ea=s.address) for s in analysis.struct]

    @mcp.tool
    def struct_dump(analysis_id:Annotated[int, "analysis ID"], ea:Annotated[Union[int,str], "effective address of structure"], max_depth:Annotated[int, "max struct depth"]=3) -> Dict[str, Any]:
        """
        dump a structure in json
        """
        analysis = mcp.get_analysis(analysis_id)
        res = None

        def field2json(sa, lvl=0, max_depth=2):
            res = None
            if sa.field.type in (malcat.Field.Type.Struct, malcat.Field.Type.Array):
                if lvl >= max_depth:
                    res = "<max depth reached>"
                else:
                    res = {}
                    if sa.field.type == malcat.Field.Type.Struct:
                        max_children = 64
                    else:
                        max_children = 16
                    for i in range(min(sa.count, max_children)):
                        sa2 = sa.at(i)
                        res[sa2.name] = field2json(sa2, lvl+1, max_depth)
                    if sa.count >= max_children:
                        res["<output truncated>"] = f"{sa.count - max_children} more fields/elements"
            elif sa.field.type == malcat.Field.Type.BitVector:
                res = []
                for i in range(sa.count - 1, -1, -1):
                    bit = sa.at(i)
                    if bit.value:
                        res.append(bit.name)
            elif isinstance(sa.value, datetime.datetime):
                res = sa.value.strftime("%Y-%m-%d %H:%M:%S")
            elif sa.field.type == malcat.Field.Type.Bytes:
                res = "{}... ({:d} bytes)\n".format(sa.value[:16].hex(), len(sa))
            else:
                res = sa.value
            return res
        
        ea = smart_ea_convert(analysis, ea)
        sa = analysis.struct.find(ea)
        if sa is None:
            raise McpToolError("No structure at address")
        return field2json(sa)

    class YaraMatch(TypedDict):
        id: Annotated[str, "Yara rule id"]
        category: Annotated[str, "Yara rule category"]
        reliability: Annotated[int, "0 (very FP-prone) to 100 (no FP at all)"]
        type: Annotated[str, "dangerosity of the rule"]
        description: Annotated[str, "what the rule detects"]
        num_patterns: Annotated[int, "number of matching string/byte patterns"]


    ################################# SYMBOLS

    class SymbolInfo(TypedDict):
        ea: Annotated[int, "effective address of symbol"]
        name: Annotated[str, "symbol name"]
        type: Annotated[str, 'symbol type. either "IMPORT", "EXPORT", "TYPE" or "DEBUG"']
        num_refs: Annotated[int, "number of timesincoming references to the symbol"]

    @mcp.tool
    def symbols_search(analysis_id:Annotated[int, "analysis ID"], 
            maximum_number_of_symbols:Annotated[int, "maximum number of symbols to return"]=128, 
            symbol_type:Annotated[str, 'symbol type. either "IMPORT", "EXPORT", "TYPE" or "DEBUG"']="", 
            contains:Annotated[str,"substring which must be found in the symbol name"]="") -> List[SymbolInfo]:
        """
        search among all symbols respecting _all_ the given search criterions. leave criterions empty to list all symbols.
        """
        analysis = mcp.get_analysis(analysis_id)
        matching = []
        type2nice = {
            malcat.Symbol.Type.IMPORT: "IMPORT",
            malcat.Symbol.Type.EXPORT: "EXPORT",
            malcat.Symbol.Type.FUNCTION: "DEBUG",
            malcat.Symbol.Type.TYPEDEF: "TYPE",
        }
        nice2type = {v: k for k,v in type2nice.items()}
        st = nice2type.get(symbol_type)
        for s in analysis.syms:
            if contains and contains not in s.name:
                continue
            if st and s.type != st:
                continue
            tt = type2nice.get(s.type)
            if tt:
                matching.append(SymbolInfo(ea=s.address, name=s.name, type=tt, num_refs=len(analysis.xref[s.address])))
        return matching

   
    ################################# PICTURES

    @mcp.tool
    def image_extract_pixels(analysis_id:Annotated[int, "analysis ID"], output_path:Annotated[str, "where to write pixel data on disk"]):
        """
        Extract (uncompressed) pixels data from a BMP or PNG file
        """
        analysis = mcp.get_analysis(analysis_id)

        if analysis.type == "PNG":
            import zlib
            hdr = analysis.struct["IHDR"]
            if hdr["Compression"] != 0:
                raise McpToolError("Unsupported compression method")
            obj = zlib.decompressobj()
            done = []
            for s in analysis.struct:
                if s.name == "IDAT":
                    done.append(obj.decompress(s["Data"]))
            data = b"".join(done) + obj.flush()
        elif "BitmapInfoHeader" in analysis.struct:
            bih = analysis.struct["BitmapInfoHeader"]
            w = bih["biWidth"]
            h = bih["biHeight"]
            bpp = bih["biBitCount"]
            data = analysis.struct["biImageData"]
            aligned_line_width = ((w * bpp + 31) // 32) * 4
            effective_line_width = (w * bpp) // 8
            res = bytearray(effective_line_width * h)
            for y in range(h):
                res[y*effective_line_width:(y+1)*effective_line_width] = data[((h-1)-y) * aligned_line_width: ((h-1)-y) * aligned_line_width + effective_line_width]
            data = bytes(res)
        else:
            raise McpToolError(f"Unsupported file type { analysis.type}")

        dirname = os.path.dirname(output_path)
        if not os.path.exists(dirname):
            os.makedirs(dirname)
        with open(output_path, "wb") as f:
            f.write(data)

    @mcp.tool
    def unpack_donut(analysis_id:Annotated[int, "analysis ID of the donut shellcode/loader"]) -> MalcatAnalysisResult:
        """
        Unpack the payload of a Donut shellcode instance and analyse it
        """
        from scripts.unpackers.donut import DonutDecryptorForMalcat
        analysis = mcp.get_analysis(analysis_id)
        instances = DonutDecryptorForMalcat.find_donuts(analysis.file)
        if not instances:
            raise McpToolError("No known DonutLoader instance found")
        last_error = ""
        for instance in instances:
            try:
                instance_data, module = instance.unpack()
                mf = malcat.File.from_buffer(module, "donut_payload")
                return MalcatAnalysisResult.from_analysis(mcp.analyse_file(mf))
            except Exception as e:
                last_error = str(e)
        else:
            raise McpToolError("Could not unpack donut: " + last_error)
