malwarebazaar_mcp.py

#!/usr/bin/env python3

"""
MalwareBazaar MCP Server

This script defines a local MCP (Model Context Protocol) server using FastMCP to query and analyze
threat intelligence data from MalwareBazaar (https://bazaar.abuse.ch/). It provides tools to retrieve
recent samples, fetch detailed metadata, download files, and query by tag, using the MalwareBazaar API.

The script is intended to be run as an MCP tool within a larger agent system and communicate via stdio
transport.

Environment:
    - MALWAREBAZAAR_API_KEY: API key for authenticating with MalwareBazaar.
    - DEBUG_MB (optional): Set to "1" to enable debug output.

Usage:
    $ uv run malwarebazaar_mcp.py

Returns:
    Formatted results or error messages based on the executed subcommand.
"""

from typing import Any
import os
import sys
import json
import base64
import pprint
import httpx
import re
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP

# Initialize FastMCP server
mcp = FastMCP("MalwareBazaar MCP")

# Constants
MBAZ_API_URL = "https://mb-api.abuse.ch/api/v1/"
USER_AGENT = "malwarebazaar-mcp/1.5"

# Load API key from environment
load_dotenv()
RAW_KEY = os.getenv("MALWAREBAZAAR_API_KEY")
API_KEY = RAW_KEY.strip() if RAW_KEY else None
DEBUG = os.getenv("DEBUG_MB") == "1"

# Ensure the API key is set before making any requests
if not API_KEY:
    raise RuntimeError("MALWAREBAZAAR_API_KEY is missing (.env or config env block)")


async def make_mb_request(payload: dict[str, Any]) -> dict[str, Any] | None:
    """
    Sends a POST request to the MalwareBazaar API with the specified payload.

    Parameters:
        payload (dict): The dictionary of form fields to send with the request.

    Returns:
        dict | None: Parsed JSON response from the API if successful, None otherwise.
    """
    headers = {
        "User-Agent": USER_AGENT,
        "Accept": "application/json",
        "Auth-Key": API_KEY,
    }

    async with httpx.AsyncClient(timeout=30.0) as client:
        try:
            request = client.build_request(
                "POST", MBAZ_API_URL, headers=headers, data=payload
            )
            if DEBUG:
                print("\n>>> --- MalwareBazaar request ---", file=sys.stderr)
                pprint.pprint(dict(request.headers), stream=sys.stderr, width=120)
                preview = base64.b64encode(request.content[:200]).decode()
                print(">>> body preview (b64):", preview, file=sys.stderr)

            response = await client.send(request)
            raw = await response.aread()

            print(
                ">>> RAW RESPONSE:",
                raw.decode(errors="replace")[:1000],
                file=sys.stderr,
            )

            if DEBUG:
                print(">>> status :", response.status_code, file=sys.stderr)
                print(">>> headers:", dict(response.headers), file=sys.stderr)

            response.raise_for_status()
            try:
                return json.loads(raw)
            except json.JSONDecodeError as e:
                print(
                    f"JSON decode failed: {e}\nRaw content: {raw[:200]}",
                    file=sys.stderr,
                )
                return None
        except Exception as e:
            print(
                f"MalwareBazaar request failed: {e}\nRequest: {payload}",
                file=sys.stderr,
            )
            return None


def format_detailed(sample: dict) -> str:
    """
    Formats and returns a multiline string of key metadata fields from a malware sample,
    prioritizing a fixed key order and appending any additional fields not recognized.

    Parameters:
        sample (dict): Dictionary containing the sample metadata.

    Returns:
        str: Formatted multiline text block.
    """
    keys_order = [...]  # fill in the desired field order
    lines = []
    for key in keys_order:
        if key in sample:
            value = sample[key]
            value_str = (
                json.dumps(value, indent=2)
                if isinstance(value, (dict, list))
                else str(value)
            )
            lines.append(f"{key}: {value_str}")
    extra_keys = set(sample.keys()) - set(keys_order)
    if extra_keys:
        lines.append("\n-- Additional Fields --")
        for key in sorted(extra_keys):
            value = sample[key]
            value_str = (
                json.dumps(value, indent=2)
                if isinstance(value, (dict, list))
                else str(value)
            )
            lines.append(f"{key}: {value_str}")
    return "\n".join(lines)


def format_basic(sample: dict) -> str:
    """
    Formats and returns a one-line summary string with hash and timestamp.

    Parameters:
        sample (dict): Dictionary containing the sample metadata.

    Returns:
        str: SHA256 hash and first seen timestamp.
    """
    return f"{sample.get('sha256_hash')}  {sample.get('first_seen')}"


@mcp.tool()
async def get_recent(selector: str = "time") -> str:
    """
    Retrieves up to 10 of the most recently submitted malware samples.

    Parameters:
        selector (str): Only accepted value is "time" (default).

    Returns:
        str: Formatted result of basic metadata for recent samples.
    """
    if selector != "time":
        return f"\u274c Error: `get_recent()` only supports `selector='time'`. Got: `{selector}`.\nUse `get_info()` to query by hash."

    data = await make_mb_request({"query": "get_recent", "selector": selector})
    if not data or data.get("query_status") == "no_results":
        return "No recent samples returned by MalwareBazaar."
    if data.get("query_status") != "ok":
        return f"\u274c Unexpected response: `{data.get('query_status')}`."

    samples = data["data"][:10]
    formatted = [format_basic(sample) for sample in samples]
    return (
        "### Recent Malware Samples\n"
        "Retrieved using `get_recent()`\n\n"
        "```text\n" + "\n".join(formatted) + "\n```"
    )


@mcp.tool()
async def get_info(selector: str = "", sha256: str = "") -> str:
    """
    Retrieves full metadata for a single malware sample given a SHA256, SHA1, or MD5 hash.

    Parameters:
        selector (str): Optional. Used for routing or alternative hash input.
        sha256 (str): The full hash of the sample to query.

    Returns:
        str: Full formatted metadata response from MalwareBazaar.
    """
    if not sha256 and re.fullmatch(r"[a-fA-F0-9]{64}", selector):
        sha256 = selector
    elif selector and selector != "":
        return f"\u274c Error: `selector` is not valid for `get_info()`. Did you mean to use `get_recent()`?"

    if not sha256:
        return "\u274c Error: No valid hash provided. Please pass `sha256=<value>`."

    data = await make_mb_request({"query": "get_info", "hash": sha256})
    if not data or data.get("query_status") in [
        "hash_not_found",
        "illegal_hash",
        "no_hash_provided",
    ]:
        return f"\u274c Error: `{data.get('query_status', 'Unknown error')}`."
    if data.get("query_status") != "ok":
        return f"\u274c Unexpected response: `{data.get('query_status')}`."

    sample = data["data"][0]
    return (
        f"### Malware Sample Metadata for `{sha256}`\n\n"
        "Retrieved using `get_info()`\n\n"
        "```text\n" + format_detailed(sample) + "\n```"
    )


@mcp.tool()
async def get_file(sha256: str) -> str:
    """
    Downloads a malware sample archive by SHA256 hash.

    Parameters:
        sha256 (str): SHA256 hash of the desired sample.

    Returns:
        str: Result message indicating file path, size, or error.
    """
    payload = {"query": "get_file", "sha256_hash": sha256}
    headers = {
        "User-Agent": USER_AGENT,
        "Accept": "application/octet-stream",
        "Auth-Key": API_KEY,
    }

    print(f"\n=== DEBUG: Attempting download for {sha256} ===")

    async with httpx.AsyncClient(timeout=60.0) as client:
        try:
            response = await client.post(MBAZ_API_URL, headers=headers, data=payload)
            if response.status_code == 200:
                content = response.content
                print(f"First 16 bytes: {content[:16].hex()}")
                if content.startswith(b"PK"):
                    try:
                        download_dir = os.path.abspath(os.getcwd())
                        file_path = os.path.join(download_dir, f"{sha256}.zip")
                        temp_path = f"{file_path}.tmp"
                        with open(temp_path, "wb") as f:
                            f.write(content)
                        os.rename(temp_path, file_path)
                        if os.path.exists(file_path):
                            return (
                                "### File Download Successful\n\n"
                                "```text\n"
                                f"Saved to: {file_path}\nSize: {os.path.getsize(file_path):,} bytes\nPassword: 'infected'\n"
                                "```"
                            )
                        return "File saved but verification failed"
                    except Exception as e:
                        return f"File save error: {str(e)}"
                elif b"query_status" in content:
                    try:
                        error = response.json()
                        return f"API Error: {error.get('query_status')}"
                    except:
                        return "API Error (malformed JSON)"
                return "Unexpected response format"
            return f"HTTP Error {response.status_code}"
        except Exception as e:
            return f"Request failed: {str(e)}"


@mcp.tool()
async def get_taginfo(tag: str, limit: int = 100) -> str:
    """
    Retrieves a list of malware samples from MalwareBazaar tagged with a specific keyword.

    Parameters:
        tag (str): The keyword tag to filter malware samples by.
        limit (int): Maximum number of results to retrieve (default 100, max 1000).

    Returns:
        str: A formatted list of results or an error string.
    """
    if limit > 1000:
        limit = 1000
    payload = {"query": "get_taginfo", "tag": tag, "limit": limit}
    data = await make_mb_request(payload)
    if not data:
        return "\u274c Error: No data received from MalwareBazaar."
    if data.get("query_status") != "ok":
        return f"\u274c Error: `{data.get('query_status', 'Unknown error')}`."
    samples = data.get("data", [])
    if not samples:
        return f"No malware samples found for tag `{tag}`."

    formatted = [format_basic(sample) for sample in samples]
    return (
        f"### Samples for Tag `{tag}`\n\n"
        f"Retrieved using `get_taginfo()` with limit `{limit}`.\n\n"
        "```text\n" + "\n".join(formatted) + "\n```"
    )


if __name__ == "__main__":
    mcp.run(transport="stdio")