iflow-mcp_developermode-korea_reversecore-mcp 1.0.0__py3-none-any.whl

reversecore_mcp/core/r2_pool.py

@@ -0,0 +1,403 @@
+"""
+Radare2 Connection Pool
+
+This module provides a connection pool for managing persistent r2pipe instances.
+It helps reduce the overhead of spawning new radare2 processes for every command.
+
+Features:
+- Configurable pool size via REVERSECORE_R2_POOL_SIZE
+- Configurable timeout via REVERSECORE_R2_POOL_TIMEOUT
+- LRU eviction policy to limit memory usage
+- Thread-safe and async-safe execution
+- Automatic reconnection on failure
+- Health checking for stale connections
+"""
+
+import asyncio
+import threading
+import time
+from collections import OrderedDict
+from collections.abc import AsyncGenerator, Generator
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any
+
+try:
+    import r2pipe
+except ImportError:
+    r2pipe = None
+
+from reversecore_mcp.core.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+class R2PoolTimeout(Exception):
+    """Raised when connection acquisition times out."""
+
+    pass
+
+
+class R2ConnectionPool:
+    """
+    Manages a pool of persistent r2pipe connections.
+
+    Features:
+    - Configurable pool size (default from config or 10)
+    - Configurable acquisition timeout
+    - LRU eviction policy to limit memory usage
+    - Thread-safe and async-safe execution
+    - Automatic reconnection on failure
+    - Health checking for stale connections
+
+    Configuration:
+        Pool size and timeout can be configured via environment variables:
+        - REVERSECORE_R2_POOL_SIZE: Maximum number of connections (default: 3)
+        - REVERSECORE_R2_POOL_TIMEOUT: Acquisition timeout in seconds (default: 30)
+    """
+
+    def __init__(
+        self,
+        max_connections: int | None = None,
+        acquisition_timeout: int | None = None,
+        health_check_interval: int = 60,
+    ):
+        """Initialize the connection pool.
+
+        Args:
+            max_connections: Maximum number of connections to maintain.
+                           If None, uses config value or default of 10.
+            acquisition_timeout: Timeout in seconds for acquiring a connection.
+                               If None, uses config value or default of 30.
+            health_check_interval: Interval in seconds for health checks.
+        """
+        # Lazy-load config to avoid circular imports
+        self._max_connections = max_connections
+        self._acquisition_timeout = acquisition_timeout
+        self._health_check_interval = health_check_interval
+        self._config_loaded = False
+
+        self._pool: OrderedDict[str, Any] = OrderedDict()
+        self._lock = threading.RLock()
+        self._async_lock: asyncio.Lock | None = None
+        self._async_lock_init_lock = threading.Lock()
+        self._last_access: dict[str, float] = {}
+        self._analyzed_files: set[str] = set()
+        self._last_health_check: dict[str, float] = {}
+
+        # Semaphore for limiting concurrent connections
+        self._connection_semaphore: threading.Semaphore | None = None
+        self._async_semaphore: asyncio.Semaphore | None = None
+
+        # Statistics
+        self._stats = {
+            "connections_created": 0,
+            "connections_evicted": 0,
+            "cache_hits": 0,
+            "cache_misses": 0,
+            "reconnections": 0,
+        }
+
+    def _load_config(self) -> None:
+        """Load configuration values lazily."""
+        if self._config_loaded:
+            return
+
+        try:
+            from reversecore_mcp.core.config import get_config
+
+            config = get_config()
+            if self._max_connections is None:
+                self._max_connections = config.r2_pool_size
+            if self._acquisition_timeout is None:
+                self._acquisition_timeout = config.r2_pool_timeout
+        except Exception:
+            # Fallback to defaults if config is not available
+            if self._max_connections is None:
+                self._max_connections = 10
+            if self._acquisition_timeout is None:
+                self._acquisition_timeout = 30
+
+        self._config_loaded = True
+
+    @property
+    def max_connections(self) -> int:
+        """Get maximum connections, loading from config if needed."""
+        self._load_config()
+        return self._max_connections or 10
+
+    @property
+    def acquisition_timeout(self) -> int:
+        """Get acquisition timeout, loading from config if needed."""
+        self._load_config()
+        return self._acquisition_timeout or 30
+
+    def _get_connection_semaphore(self) -> threading.Semaphore:
+        """Get or create a connection semaphore for rate limiting."""
+        if self._connection_semaphore is None:
+            self._connection_semaphore = threading.Semaphore(self.max_connections)
+        return self._connection_semaphore
+
+    def _get_async_semaphore(self) -> asyncio.Semaphore:
+        """Get or create an async semaphore for rate limiting."""
+        if self._async_semaphore is None:
+            self._async_semaphore = asyncio.Semaphore(self.max_connections)
+        return self._async_semaphore
+
+    def _get_async_lock(self) -> asyncio.Lock:
+        """Get or create an async lock for thread-safe async operations.
+
+        The lock is lazily initialized to ensure it's created in the correct
+        event loop context. Uses double-checked locking pattern for thread-safety.
+        """
+        if self._async_lock is not None:
+            return self._async_lock
+
+        with self._async_lock_init_lock:
+            if self._async_lock is None:
+                self._async_lock = asyncio.Lock()
+            return self._async_lock
+
+    def _is_connection_healthy(self, file_path: str, r2: Any) -> bool:
+        """Check if a connection is still healthy."""
+        try:
+            # Quick health check: try to get current seek position
+            result = r2.cmd("s")
+            return result is not None
+        except Exception:
+            return False
+
+    def _maybe_health_check(self, file_path: str, r2: Any) -> bool:
+        """Perform health check if enough time has passed."""
+        now = time.time()
+        last_check = self._last_health_check.get(file_path, 0)
+
+        if now - last_check > self._health_check_interval:
+            self._last_health_check[file_path] = now
+            return self._is_connection_healthy(file_path, r2)
+        return True  # Assume healthy if recently checked
+
+    def get_connection(self, file_path: str) -> Any:
+        """Get or create an r2pipe connection for the given file."""
+        if r2pipe is None:
+            raise ImportError("r2pipe is not installed")
+
+        with self._lock:
+            self._last_access[file_path] = time.time()
+
+            if file_path in self._pool:
+                r2 = self._pool[file_path]
+
+                # Health check
+                if not self._maybe_health_check(file_path, r2):
+                    logger.warning(f"Stale connection for {file_path}, reconnecting")
+                    self._remove_connection_unsafe(file_path)
+                else:
+                    self._pool.move_to_end(file_path)
+                    self._stats["cache_hits"] += 1
+                    return r2
+
+            self._stats["cache_misses"] += 1
+
+            # Evict if full
+            while len(self._pool) >= self.max_connections:
+                oldest_file, oldest_r2 = self._pool.popitem(last=False)
+                logger.debug(f"Evicting r2 connection for {oldest_file}")
+                self._stats["connections_evicted"] += 1
+                try:
+                    oldest_r2.quit()
+                except Exception as e:
+                    logger.warning(f"Error closing r2 connection: {e}")
+                self._last_access.pop(oldest_file, None)
+                self._last_health_check.pop(oldest_file, None)
+                self._analyzed_files.discard(oldest_file)
+
+            # Create new connection
+            logger.info(f"Opening new r2 connection for {file_path}")
+            try:
+                r2 = r2pipe.open(file_path, flags=["-2"])
+                self._pool[file_path] = r2
+                self._last_health_check[file_path] = time.time()
+                self._stats["connections_created"] += 1
+                return r2
+            except Exception as e:
+                logger.error(f"Failed to open r2 connection for {file_path}: {e}")
+                raise
+
+    def _remove_connection_unsafe(self, file_path: str) -> None:
+        """Remove a connection without locking (caller must hold lock)."""
+        if file_path in self._pool:
+            try:
+                self._pool[file_path].quit()
+            except Exception:
+                pass
+            del self._pool[file_path]
+        self._last_access.pop(file_path, None)
+        self._last_health_check.pop(file_path, None)
+        self._analyzed_files.discard(file_path)
+
+    def execute(self, file_path: str, command: str) -> str:
+        """Execute a command on the r2 connection for the given file."""
+        with self._lock:
+            try:
+                r2 = self.get_connection(file_path)
+                return r2.cmd(command)
+            except Exception as e:
+                logger.warning(f"r2 command failed, retrying connection: {e}")
+                self._remove_connection_unsafe(file_path)
+                self._stats["reconnections"] += 1
+
+                try:
+                    r2 = self.get_connection(file_path)
+                    return r2.cmd(command)
+                except Exception as retry_error:
+                    logger.error(f"Retry failed: {retry_error}")
+                    raise
+
+    async def execute_async(self, file_path: str, command: str) -> str:
+        """Execute a command asynchronously with proper async lock."""
+        async with self._get_async_lock():
+            return await asyncio.to_thread(self._execute_unsafe, file_path, command)
+
+    def _execute_unsafe(self, file_path: str, command: str) -> str:
+        """Execute with thread lock for safe asyncio.to_thread usage.
+        
+        Note: Despite the name 'unsafe', this method now acquires self._lock
+        to ensure thread-safety when called from asyncio.to_thread().
+        The async lock in execute_async() serializes async callers,
+        while this thread lock protects against concurrent sync callers.
+        """
+        with self._lock:  # Thread lock for safe pool access
+            try:
+                r2 = self._get_connection_unsafe(file_path)
+                return r2.cmd(command)
+            except Exception as e:
+                logger.warning(f"r2 command failed, retrying connection: {e}")
+                self._remove_connection_unsafe(file_path)
+                self._stats["reconnections"] += 1
+
+                try:
+                    r2 = self._get_connection_unsafe(file_path)
+                    return r2.cmd(command)
+                except Exception as retry_error:
+                    logger.error(f"Retry failed: {retry_error}")
+                    raise
+
+    def _get_connection_unsafe(self, file_path: str) -> Any:
+        """Get or create connection without locking (caller must hold lock)."""
+        if r2pipe is None:
+            raise ImportError("r2pipe is not installed")
+
+        self._last_access[file_path] = time.time()
+
+        if file_path in self._pool:
+            r2 = self._pool[file_path]
+
+            if not self._maybe_health_check(file_path, r2):
+                logger.warning(f"Stale connection for {file_path}, reconnecting")
+                self._remove_connection_unsafe(file_path)
+            else:
+                self._pool.move_to_end(file_path)
+                self._stats["cache_hits"] += 1
+                return r2
+
+        self._stats["cache_misses"] += 1
+
+        # Evict if full
+        while len(self._pool) >= self.max_connections:
+            oldest_file, oldest_r2 = self._pool.popitem(last=False)
+            logger.debug(f"Evicting r2 connection for {oldest_file}")
+            self._stats["connections_evicted"] += 1
+            try:
+                oldest_r2.quit()
+            except Exception as e:
+                logger.warning(f"Error closing r2 connection: {e}")
+            self._last_access.pop(oldest_file, None)
+            self._last_health_check.pop(oldest_file, None)
+            self._analyzed_files.discard(oldest_file)
+
+        # Create new connection
+        logger.info(f"Opening new r2 connection for {file_path}")
+        try:
+            r2 = r2pipe.open(file_path, flags=["-2"])
+            self._pool[file_path] = r2
+            self._last_health_check[file_path] = time.time()
+            self._stats["connections_created"] += 1
+            return r2
+        except Exception as e:
+            logger.error(f"Failed to open r2 connection for {file_path}: {e}")
+            raise
+
+    @asynccontextmanager
+    async def async_session(self, file_path: str) -> AsyncGenerator[Any, None]:
+        """Async context manager for r2 connection.
+
+        Usage:
+            async with r2_pool.async_session(path) as r2:
+                result = r2.cmd('aaa')
+        """
+        async with self._get_async_lock():
+            r2 = await asyncio.to_thread(self._get_connection_unsafe, file_path)
+            try:
+                yield r2
+            except Exception as e:
+                logger.warning(f"Error in async session: {e}")
+                # Invalidate connection on error
+                if file_path in self._pool:
+                    del self._pool[file_path]
+                raise
+
+    @contextmanager
+    def sync_session(self, file_path: str) -> Generator[Any, None, None]:
+        """Sync context manager for r2 connection.
+
+        Usage:
+            with r2_pool.sync_session(path) as r2:
+                result = r2.cmd('aaa')
+        """
+        with self._lock:
+            r2 = self._get_connection_unsafe(file_path)
+            try:
+                yield r2
+            except Exception as e:
+                logger.warning(f"Error in sync session: {e}")
+                self._remove_connection_unsafe(file_path)
+                raise
+
+    def close_all(self):
+        """Close all connections in the pool."""
+        with self._lock:
+            for _file_path, r2 in self._pool.items():
+                try:
+                    r2.quit()
+                except Exception:
+                    pass
+            self._pool.clear()
+            self._last_access.clear()
+            self._last_health_check.clear()
+            self._analyzed_files.clear()
+
+    def is_analyzed(self, file_path: str) -> bool:
+        """Check if the file has been analyzed."""
+        with self._lock:
+            return file_path in self._analyzed_files
+
+    def mark_analyzed(self, file_path: str):
+        """Mark the file as analyzed."""
+        with self._lock:
+            if file_path in self._pool:
+                self._analyzed_files.add(file_path)
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get pool statistics."""
+        with self._lock:
+            return {
+                **self._stats,
+                "current_connections": len(self._pool),
+                "max_connections": self.max_connections,
+                "analyzed_files": len(self._analyzed_files),
+            }
+
+
+# Global instance (for backward compatibility)
+# New code should use: from reversecore_mcp.core.container import get_r2_pool
+r2_pool = R2ConnectionPool()

reversecore_mcp/core/report_generator.py

@@ -0,0 +1,268 @@
+"""
+Evidence-based Report Generator for Malware Analysis.
+
+This module generates professional SOC/IR reports with evidence tracking,
+confidence levels, and clear differentiation between observed facts and inferences.
+"""
+
+from datetime import datetime
+from typing import Any, Optional
+from dataclasses import dataclass, field
+
+from reversecore_mcp.core.evidence import (
+    EvidenceLevel,
+    MITREConfidence,
+    Finding,
+    Evidence,
+    MITRETechnique,
+    AnalysisMetadata,
+)
+from reversecore_mcp.core import json_utils as json
+
+
+@dataclass
+class EvidenceBasedReport:
+    """Evidence-based malware analysis report."""
+    
+    metadata: AnalysisMetadata
+    executive_summary: str = ""
+    malware_family: str = "Unknown"
+    family_confidence: float = 0.0
+    family_evidence: list[str] = field(default_factory=list)
+    
+    findings: list[Finding] = field(default_factory=list)
+    mitre_techniques: list[MITRETechnique] = field(default_factory=list)
+    iocs: dict[str, list[str]] = field(default_factory=dict)
+    
+    recommendations: list[str] = field(default_factory=list)
+    yara_rule: Optional[str] = None
+    
+    def add_finding(self, finding: Finding) -> None:
+        """Add a finding to the report."""
+        self.findings.append(finding)
+    
+    def add_mitre(self, technique: MITRETechnique) -> None:
+        """Add a MITRE technique mapping."""
+        self.mitre_techniques.append(technique)
+    
+    def add_ioc(self, ioc_type: str, value: str) -> None:
+        """Add an IOC."""
+        if ioc_type not in self.iocs:
+            self.iocs[ioc_type] = []
+        if value not in self.iocs[ioc_type]:
+            self.iocs[ioc_type].append(value)
+    
+    def set_family(self, family: str, confidence: float, evidence: list[str]) -> None:
+        """Set malware family with confidence and evidence."""
+        self.malware_family = family
+        self.family_confidence = confidence
+        self.family_evidence = evidence
+    
+    def finalize(self) -> None:
+        """Finalize the report (set end time)."""
+        self.metadata.end_time = datetime.now()
+    
+    @property
+    def observed_count(self) -> int:
+        return sum(1 for f in self.findings if f.level == EvidenceLevel.OBSERVED)
+    
+    @property
+    def inferred_count(self) -> int:
+        return sum(1 for f in self.findings if f.level == EvidenceLevel.INFERRED)
+    
+    @property
+    def possible_count(self) -> int:
+        return sum(1 for f in self.findings if f.level == EvidenceLevel.POSSIBLE)
+    
+    @property
+    def overall_confidence(self) -> float:
+        """Calculate overall report confidence."""
+        if not self.findings:
+            return 0.0
+        return sum(f.confidence for f in self.findings) / len(self.findings)
+    
+    def generate_markdown(self) -> str:
+        """Generate professional markdown report."""
+        lines = []
+        
+        # Header
+        lines.append(f"# 🔬 Malware Analysis Report")
+        lines.append("")
+        lines.append(f"**Sample**: {self.metadata.sample_name}")
+        lines.append(f"**SHA256**: `{self.metadata.sample_hash}`")
+        lines.append(f"**Analysis Date**: {self.metadata.start_time.strftime('%Y-%m-%d %H:%M:%S')}")
+        lines.append(f"**Duration**: {self.metadata.duration_formatted}")
+        lines.append(f"**Analyst**: {self.metadata.analyst}")
+        lines.append("")
+        
+        # Confidence Summary Box
+        lines.append("---")
+        lines.append("")
+        lines.append("## 📊 Confidence Summary")
+        lines.append("")
+        lines.append("| Metric | Value |")
+        lines.append("|--------|-------|")
+        lines.append(f"| **Overall Confidence** | {self.overall_confidence:.0%} |")
+        lines.append(f"| 🔍 Observed Findings | {self.observed_count} |")
+        lines.append(f"| 🔎 Inferred Findings | {self.inferred_count} |")
+        lines.append(f"| ❓ Possible Findings | {self.possible_count} |")
+        lines.append("")
+        
+        # Malware Family Identification
+        lines.append("---")
+        lines.append("")
+        lines.append("## 🦠 Malware Identification")
+        lines.append("")
+        
+        if self.family_confidence >= 0.8:
+            verdict = "✅ **CONFIRMED**"
+        elif self.family_confidence >= 0.6:
+            verdict = "🟡 **LIKELY**"
+        elif self.family_confidence >= 0.4:
+            verdict = "🟠 **POSSIBLE**"
+        else:
+            verdict = "❓ **UNCERTAIN**"
+        
+        lines.append(f"**Family**: {self.malware_family}")
+        lines.append(f"**Confidence**: {self.family_confidence:.0%} {verdict}")
+        lines.append("")
+        
+        if self.family_evidence:
+            lines.append("**Identification Evidence:**")
+            for ev in self.family_evidence:
+                lines.append(f"  - {ev}")
+            lines.append("")
+        
+        # Executive Summary
+        if self.executive_summary:
+            lines.append("---")
+            lines.append("")
+            lines.append("## 📋 Executive Summary")
+            lines.append("")
+            lines.append(self.executive_summary)
+            lines.append("")
+        
+        # Findings by Evidence Level
+        lines.append("---")
+        lines.append("")
+        lines.append("## 🔍 Analysis Findings")
+        lines.append("")
+        lines.append("> **Legend**: 🔍 Observed (verified) | 🔎 Inferred (high confidence) | ❓ Possible (needs verification)")
+        lines.append("")
+        
+        # Group findings by level
+        for level in [EvidenceLevel.OBSERVED, EvidenceLevel.INFERRED, EvidenceLevel.POSSIBLE]:
+            level_findings = [f for f in self.findings if f.level == level]
+            if level_findings:
+                lines.append(f"### {level.symbol} {level.value.upper()} Findings ({len(level_findings)})")
+                lines.append("")
+                for finding in level_findings:
+                    lines.append(finding.format_markdown())
+                    lines.append("")
+        
+        # MITRE ATT&CK Mapping
+        if self.mitre_techniques:
+            lines.append("---")
+            lines.append("")
+            lines.append("## ⚔️ MITRE ATT&CK Mapping")
+            lines.append("")
+            lines.append("> **Confidence Levels**: ✅ Confirmed | 🟢 High | 🟡 Medium | 🔴 Low")
+            lines.append("")
+            lines.append("| Technique ID | Name | Tactic | Confidence |")
+            lines.append("|-------------|------|--------|------------|")
+            for tech in self.mitre_techniques:
+                lines.append(tech.format_markdown_row())
+            lines.append("")
+        
+        # IOCs
+        if self.iocs:
+            lines.append("---")
+            lines.append("")
+            lines.append("## 🎯 Indicators of Compromise (IOCs)")
+            lines.append("")
+            for ioc_type, values in self.iocs.items():
+                lines.append(f"### {ioc_type.title()} ({len(values)})")
+                lines.append("")
+                for val in values:
+                    lines.append(f"- `{val}`")
+                lines.append("")
+        
+        # YARA Rule
+        if self.yara_rule:
+            lines.append("---")
+            lines.append("")
+            lines.append("## 📝 Detection Rule (YARA)")
+            lines.append("")
+            lines.append("```yara")
+            lines.append(self.yara_rule)
+            lines.append("```")
+            lines.append("")
+        
+        # Recommendations
+        if self.recommendations:
+            lines.append("---")
+            lines.append("")
+            lines.append("## 🛡️ Recommendations")
+            lines.append("")
+            for i, rec in enumerate(self.recommendations, 1):
+                lines.append(f"{i}. {rec}")
+            lines.append("")
+        
+        # Footer
+        lines.append("---")
+        lines.append("")
+        lines.append("## 📎 Report Metadata")
+        lines.append("")
+        lines.append(f"- **Session ID**: {self.metadata.session_id}")
+        lines.append(f"- **Start Time**: {self.metadata.start_time.isoformat()}")
+        if self.metadata.end_time:
+            lines.append(f"- **End Time**: {self.metadata.end_time.isoformat()}")
+        lines.append(f"- **Duration**: {self.metadata.duration_formatted}")
+        lines.append(f"- **Tools Used**: {', '.join(self.metadata.tools_used) if self.metadata.tools_used else 'N/A'}")
+        lines.append("")
+        lines.append("---")
+        lines.append(f"*Generated by Reversecore MCP Server*")
+        
+        return "\n".join(lines)
+    
+    def to_dict(self) -> dict[str, Any]:
+        """Convert report to dictionary for JSON export."""
+        return {
+            "metadata": self.metadata.to_dict(),
+            "executive_summary": self.executive_summary,
+            "malware_family": self.malware_family,
+            "family_confidence": self.family_confidence,
+            "family_evidence": self.family_evidence,
+            "confidence_summary": {
+                "overall": round(self.overall_confidence, 2),
+                "observed_count": self.observed_count,
+                "inferred_count": self.inferred_count,
+                "possible_count": self.possible_count,
+            },
+            "findings": [f.to_dict() for f in self.findings],
+            "mitre_techniques": [t.to_dict() for t in self.mitre_techniques],
+            "iocs": self.iocs,
+            "recommendations": self.recommendations,
+            "yara_rule": self.yara_rule,
+        }
+    
+    def to_json(self, indent: int = 2) -> str:
+        """Export report as JSON."""
+        return json.dumps(self.to_dict(), indent=indent)
+
+
+def create_report(
+    session_id: str,
+    sample_name: str,
+    sample_hash: str,
+    analyst: str = "Reversecore MCP",
+) -> EvidenceBasedReport:
+    """Create a new evidence-based report."""
+    metadata = AnalysisMetadata(
+        session_id=session_id,
+        sample_name=sample_name,
+        sample_hash=sample_hash,
+        start_time=datetime.now(),
+        analyst=analyst,
+    )
+    return EvidenceBasedReport(metadata=metadata)