agentic-threat-hunting-framework 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

{agentic_threat_hunting_framework-0.4.0.dist-info → agentic_threat_hunting_framework-0.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentic-threat-hunting-framework
-Version: 0.4.0
+Version: 0.5.0
 Summary: Agentic Threat Hunting Framework - Memory and AI for threat hunters
 Author-email: Sydney Marrone <athf@nebulock.io>
 Maintainer-email: Sydney Marrone <athf@nebulock.io>

{agentic_threat_hunting_framework-0.4.0.dist-info → agentic_threat_hunting_framework-0.5.0.dist-info}/RECORD

@@ -1,13 +1,14 @@
-agentic_threat_hunting_framework-0.4.0.dist-info/licenses/LICENSE,sha256=_KObErRfiKoolznt-DF0nJnr3U9Rdh7Z4Ba7G5qqckk,1071
+agentic_threat_hunting_framework-0.5.0.dist-info/licenses/LICENSE,sha256=_KObErRfiKoolznt-DF0nJnr3U9Rdh7Z4Ba7G5qqckk,1071
 athf/__init__.py,sha256=OrjZe8P97_BTEkscapnwSsqKSjwXNP9d8-HtGr19Ni0,241
 athf/__version__.py,sha256=wCIQoU9b7qKcSNQiIOgHaD2buzBC-dlQYtvg8X5WS4A,59
-athf/cli.py,sha256=l8PkZCs4ToUOAxX_Ciy5bLD1mV1Qta3_WalBJOA1t9c,4787
+athf/cli.py,sha256=108PnDRlaytJj9KzjzcTLljB3DeerMIXZOeAQJrmtPU,5052
+athf/plugin_system.py,sha256=c_9-oUiR6tuYpWpEmeVRayU8-TXlkjvZC3EUxuYWW4M,1515
 athf/agents/__init__.py,sha256=iaSJpvnXm9rz4QS7gBrsaLEjm49uvsMs4BLPOJeyp78,346
-athf/agents/base.py,sha256=HQJpQAKmY7folbzPnvF2OZ9vhEf205lBFUHrtfOXr64,4352
+athf/agents/base.py,sha256=lnVDIOQUOyP-Apa9UM2E1mRXUPnNJ4hVqQXOwVw2u4c,4286
 athf/agents/llm/__init__.py,sha256=qSGA-NaInjsDkMpGQwnTz3S1OgCVlzetpMcDS_to1co,671
 athf/agents/llm/hunt_researcher.py,sha256=dIyD2Izh3zdf62kCHug1DwXFgmWhOMQUTim7qM3UAIs,27071
 athf/agents/llm/hypothesis_generator.py,sha256=XkbJz8IS4zwQjEy-ZD0zy2XW5uRnAy87Lii-5XTY0WU,8564
-athf/commands/__init__.py,sha256=YYPIbajQf7DN_h8PFNfktgny4oHKKWKGbAU6tItE5vE,500
+athf/commands/__init__.py,sha256=-ZOfg6uV1eSh7RDW7dKzdufuYvQTT0KGMF4JB6waHsY,635
 athf/commands/agent.py,sha256=c7ZeZa3OArXyXTgVjmUB2JXa3m9IpLFJ_FEVDhaDLE8,19000
 athf/commands/context.py,sha256=V-at81-OgKcLY-In48-AccTnHfTgdofmnjE8S5kypoI,12678
 athf/commands/env.py,sha256=JPKRsv48cgsIAjSFaGJ1-Nu0nQKGSVg4AbiFxb9jVX4,11887
@@ -15,17 +16,24 @@ athf/commands/hunt.py,sha256=aQdgNddqy_VrxZOkxhuPxIr4KLZtX5a2ZLb9079vLlw,25169
 athf/commands/init.py,sha256=Qn0iETNyuQvM-ySqCeoDz-pPemeuzROX_karQF5yN_o,12685
 athf/commands/investigate.py,sha256=mK_id5vjfN_ukqB_-fyia0FNa0pBmtn0Xv6CKHQI1Qo,24663
 athf/commands/research.py,sha256=FrLph4agaGQ_rIxMh0OQwh1MIGDFtj40zJ3E1ZFwaAw,18112
-athf/commands/similar.py,sha256=qy1Gng73_VpqfoLLNUdxF1GBx-X29g-k8Q_wrEx26hA,11868
+athf/commands/similar.py,sha256=FTTVr4zzP9bdJrirscp6pOxdQbE8zot6pa20-_TYiuo,11804
 athf/commands/splunk.py,sha256=7n7Jl1ExqZCNxUhG0kAKgAvZMqbIoGSgx2Moq7vAu-Y,11622
 athf/core/__init__.py,sha256=yG7C8ljx3UW4QZoYvDjUxsWHlbS8M-GLGB7Je7rRfqo,31
 athf/core/attack_matrix.py,sha256=QZKKmxckQ6-U7lqVdGUJoj2jEAhP3Juvr3sqaNx2oTw,3238
+athf/core/clickhouse_connection.py,sha256=8thmJvd2pUeeRZmDE7K491NgbC0myNZsdA29ooJRfVM,13561
 athf/core/hunt_manager.py,sha256=PFsg8Ecg94NCpuFZpApo82lyORkgK5IfOIih-7-XsmM,11580
 athf/core/hunt_parser.py,sha256=FUj0yyBIcZnaS9aItMImeBDhegQwpkewIwUMNXW_ZWU,5122
 athf/core/investigation_parser.py,sha256=wbfjnq4gFgIc0a4bHIAnidVNPhbHDpIXWY1SGLk0Xls,6804
+athf/core/metrics_tracker.py,sha256=VYEiO5QVteTtR4ddyHkL61KrO4QVNUDdNaDOVFcHy4Q,18873
+athf/core/query_executor.py,sha256=OtzUkxoOdDC4ZErVIbf0Qov82uHRJ8dJ965r4pLbiVA,6271
+athf/core/query_parser.py,sha256=Uz3ZMpd4YWKLPoge16uKZLlcMQrg49Z0NLXSceg893w,6722
+athf/core/query_suggester.py,sha256=i3P624tXb9uRKGxTpcSZx4ZVbOwnCiJqLnkxQD_UqyA,7736
+athf/core/query_validator.py,sha256=mfwdtLcPZS6ON4AlR-4d8YbQ12cqpnIq6526obdPDx8,9101
 athf/core/research_manager.py,sha256=i4fUjuZJcAik8I4pwbLkQlu6cuxkWDlqaIRQrzAfB0s,14512
+athf/core/session_manager.py,sha256=8Mz082ex87VXPiSFYRFNAb9e3ED6luCy0Q6zilyaz9A,25108
 athf/core/splunk_client.py,sha256=Xib2zVwV2l8eChzqUahI3PZ7Z2XS2wz01sPbF1E0Q18,11611
 athf/core/template_engine.py,sha256=Awp0n9E5Q1dYA35XDKKAd5VJLdpaDl2N967hackUVa8,6010
-athf/core/web_search.py,sha256=lBdApqIemV2kH_NJ3vDd3adH9DwrPjaq_fs5qMjR8mI,10354
+athf/core/web_search.py,sha256=B9IhmwH7gy2RVA6WSN3L7yGp3Q4L8OsiiwcEvnnZejU,10320
 athf/data/__init__.py,sha256=QtgONloCaS3E9Ow995FMxyy6BbszpfmYeWpySQ2b9Mc,502
 athf/data/docs/CHANGELOG.md,sha256=JKkzzs1n5jSERHFi6fDt6sYEe52MSaY127dfzthkUA8,8655
 athf/data/docs/CLI_REFERENCE.md,sha256=pb76UqkY_WHJMBEXwEmK0TJR8kcGzoBPlJ0WdGMKDQM,54875
@@ -51,8 +59,8 @@ athf/data/prompts/ai-workflow.md,sha256=rZtOcGuAEi35qx7182TwHJEORdz1-RxkZMBVkg61
 athf/data/prompts/basic-prompts.md,sha256=2bunpO35RoBdJWYthXVi40RNl2UWrfwOaFthBLHF5sU,8463
 athf/data/templates/HUNT_LOCK.md,sha256=zXxHaKMWbRDLewLTegYJMbXRM72s9gFFvjdwFfGNeJE,7386
 athf/utils/__init__.py,sha256=aEAPI1xnAsowOtc036cCb9ZOek5nrrfevu8PElhbNgk,30
-agentic_threat_hunting_framework-0.4.0.dist-info/METADATA,sha256=KpwdeeDnNEGhS6zlI3NHdcKGwaN2iGYYJDIQkRr6D9E,15949
-agentic_threat_hunting_framework-0.4.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-agentic_threat_hunting_framework-0.4.0.dist-info/entry_points.txt,sha256=GopR2iTiBs-yNMWiUZ2DaFIFglXxWJx1XPjTa3ePtfE,39
-agentic_threat_hunting_framework-0.4.0.dist-info/top_level.txt,sha256=Cxxg6SMLfawDJWBITsciRzq27XV8fiaAor23o9Byoes,5
-agentic_threat_hunting_framework-0.4.0.dist-info/RECORD,,
+agentic_threat_hunting_framework-0.5.0.dist-info/METADATA,sha256=mM_lQGR-f8k7s905FXJ5xucVAoc6hp5yrq8cQmKJ-T0,15949
+agentic_threat_hunting_framework-0.5.0.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+agentic_threat_hunting_framework-0.5.0.dist-info/entry_points.txt,sha256=GopR2iTiBs-yNMWiUZ2DaFIFglXxWJx1XPjTa3ePtfE,39
+agentic_threat_hunting_framework-0.5.0.dist-info/top_level.txt,sha256=Cxxg6SMLfawDJWBITsciRzq27XV8fiaAor23o9Byoes,5
+agentic_threat_hunting_framework-0.5.0.dist-info/RECORD,,

{agentic_threat_hunting_framework-0.4.0.dist-info → agentic_threat_hunting_framework-0.5.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

athf/agents/base.py

@@ -98,7 +98,7 @@ class LLMAgent(Agent[InputT, OutputT]):
             duration_ms: Call duration in milliseconds
         """
         try:
-            from athf.core.metrics_tracker import MetricsTracker  # type: ignore[import-not-found]
+            from athf.core.metrics_tracker import MetricsTracker
 
             MetricsTracker.get_instance().log_bedrock_call(
                 agent=agent_name,
@@ -125,7 +125,7 @@ class LLMAgent(Agent[InputT, OutputT]):
             return None
 
         try:
-            import boto3  # type: ignore[import-untyped]
+            import boto3
 
             # Get AWS region from environment or use default
             region = os.getenv("AWS_REGION", os.getenv("AWS_DEFAULT_REGION", "us-east-1"))

athf/cli.py

@@ -95,8 +95,16 @@ cli.add_command(similar)
 # Agent commands
 cli.add_command(agent)
 
-# Integration commands
-cli.add_command(splunk)
+# Integration commands (optional, requires additional dependencies)
+if splunk is not None:
+    cli.add_command(splunk)
+
+# Load and register plugins
+from athf.plugin_system import PluginRegistry
+
+PluginRegistry.load_plugins()
+for name, cmd in PluginRegistry._commands.items():
+    cli.add_command(cmd, name=name)
 
 
 @cli.command(hidden=True)

athf/commands/__init__.py

@@ -7,7 +7,12 @@ from athf.commands.init import init
 from athf.commands.investigate import investigate
 from athf.commands.research import research
 from athf.commands.similar import similar
-from athf.commands.splunk import splunk
+
+# Optional: Splunk integration (requires requests package)
+try:
+    from athf.commands.splunk import splunk
+except ImportError:
+    splunk = None  # type: ignore[assignment]
 
 __all__ = [
     "init",

athf/commands/similar.py

@@ -132,8 +132,8 @@ def _find_similar_hunts(
 ) -> List[Dict[str, Any]]:
     """Find similar hunts using TF-IDF similarity."""
     try:
-        from sklearn.feature_extraction.text import TfidfVectorizer  # type: ignore[import-untyped]
-        from sklearn.metrics.pairwise import cosine_similarity  # type: ignore[import-untyped]
+        from sklearn.feature_extraction.text import TfidfVectorizer
+        from sklearn.metrics.pairwise import cosine_similarity
     except ImportError:
         console.print("[red]Error: scikit-learn not installed[/red]")
         console.print("[dim]Install with: pip install scikit-learn[/dim]")

athf/core/clickhouse_connection.py

@@ -0,0 +1,396 @@
+"""ClickHouse connection management and configuration."""
+
+import os
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import yaml
+
+
+# Custom exceptions
+class ClickHouseConfigError(Exception):
+    """Configuration validation or loading error."""
+
+    pass
+
+
+class ClickHouseConnectionError(Exception):
+    """Connection establishment or network error."""
+
+    pass
+
+
+class ClickHouseQueryError(Exception):
+    """Query execution error."""
+
+    pass
+
+
+@dataclass
+class ClickHouseConfig:
+    """Configuration for ClickHouse connection.
+
+    Attributes:
+        host: ClickHouse server hostname
+        port: ClickHouse server port (default: 8443 for HTTPS)
+        username: Database username (required)
+        password: Database password (required)
+        database: Default database name (default: "default")
+        secure: Use SSL/TLS encryption (default: True)
+    """
+
+    host: str
+    port: int
+    username: str
+    password: str
+    database: str = "default"
+    secure: bool = True
+
+    @classmethod
+    def load(cls) -> "ClickHouseConfig":
+        """Load configuration from environment variables and config file.
+
+        Precedence order:
+        1. Environment variables (highest priority)
+        2. Config file (~/.athf/clickhouse.yaml)
+        3. Hardcoded defaults (host, port, database only)
+
+        Returns:
+            ClickHouseConfig instance
+
+        Raises:
+            ClickHouseConfigError: If required credentials are missing
+        """
+        config_data: Dict[str, Any] = {}
+
+        # Load from config file first
+        config_file = Path.home() / ".athf" / "clickhouse.yaml"
+        if config_file.exists():
+            try:
+                with open(config_file, "r") as f:
+                    yaml_data = yaml.safe_load(f) or {}
+                    clickhouse_section = yaml_data.get("clickhouse", {})
+                    config_data.update(clickhouse_section)
+            except Exception as e:
+                raise ClickHouseConfigError(f"Failed to load config file {config_file}: {e}")
+
+        # Override with environment variables
+        if host := os.getenv("CLICKHOUSE_HOST"):
+            config_data["host"] = host
+        if port := os.getenv("CLICKHOUSE_PORT"):
+            config_data["port"] = int(port)
+        if user := os.getenv("CLICKHOUSE_USER"):
+            config_data["username"] = user
+        if password := os.getenv("CLICKHOUSE_PASSWORD"):
+            config_data["password"] = password
+        if database := os.getenv("CLICKHOUSE_DATABASE"):
+            config_data["database"] = database
+        if secure := os.getenv("CLICKHOUSE_SECURE"):
+            config_data["secure"] = secure.lower() in ("true", "1", "yes")
+
+        # Apply defaults
+        config_data.setdefault("host", "ohma99qewu.us-east-1.aws.clickhouse.cloud")
+        config_data.setdefault("port", 8443)
+        config_data.setdefault("database", "default")
+        config_data.setdefault("secure", True)
+
+        # Validate required fields
+        if "username" not in config_data or not config_data["username"]:
+            raise ClickHouseConfigError(
+                "Missing required credential: CLICKHOUSE_USER environment variable not set. "
+                "Set credentials with: export CLICKHOUSE_USER='your_username'"
+            )
+        if "password" not in config_data or not config_data["password"]:
+            raise ClickHouseConfigError(
+                "Missing required credential: CLICKHOUSE_PASSWORD environment variable not set. "
+                "Set credentials with: export CLICKHOUSE_PASSWORD='your_password'"
+            )
+
+        return cls(
+            host=config_data["host"],
+            port=config_data["port"],
+            username=config_data["username"],
+            password=config_data["password"],
+            database=config_data["database"],
+            secure=config_data["secure"],
+        )
+
+    def to_dict(self, mask_password: bool = True) -> Dict[str, Any]:
+        """Convert config to dictionary.
+
+        Args:
+            mask_password: If True, replace password with asterisks
+
+        Returns:
+            Dictionary representation of config
+        """
+        return {
+            "host": self.host,
+            "port": self.port,
+            "username": self.username,
+            "password": "***" if mask_password else self.password,
+            "database": self.database,
+            "secure": self.secure,
+        }
+
+
+class ClickHouseConnectionManager:
+    """Singleton connection manager for ClickHouse queries.
+
+    Manages a single ClickHouse client instance with lazy initialization.
+    Connection is reused across multiple queries within the same process.
+    """
+
+    _instance: Optional["ClickHouseConnectionManager"] = None
+    _client: Optional[Any] = None  # Type: clickhouse_connect.driver.Client
+    _config: Optional[ClickHouseConfig] = None
+
+    def __new__(cls) -> "ClickHouseConnectionManager":
+        """Ensure only one instance exists (singleton pattern)."""
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    @classmethod
+    def get_instance(cls) -> "ClickHouseConnectionManager":
+        """Get the singleton instance.
+
+        Returns:
+            ClickHouseConnectionManager instance
+        """
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    def get_client(self) -> Any:
+        """Get or create ClickHouse client (lazy initialization).
+
+        Returns:
+            ClickHouse client instance
+
+        Raises:
+            ClickHouseConnectionError: If connection fails
+        """
+        if self._client is None:
+            self._client = self._create_client()
+        return self._client
+
+    def _create_client(self) -> Any:
+        """Create ClickHouse client from configuration.
+
+        Returns:
+            ClickHouse client instance
+
+        Raises:
+            ClickHouseConnectionError: If client creation fails
+            ClickHouseConfigError: If configuration is invalid
+        """
+        try:
+            import clickhouse_connect
+        except ImportError:
+            raise ClickHouseConnectionError(
+                "clickhouse-connect not installed. Install with: pip install 'hunt-vault[clickhouse]'"
+            )
+
+        # Load configuration
+        if self._config is None:
+            self._config = ClickHouseConfig.load()
+
+        # Create client with retry logic
+        max_retries = 2
+        retry_delay = 5  # seconds
+
+        for attempt in range(max_retries):
+            try:
+                # Check if running in AWS Lambda (no SSL verification)
+                is_lambda = os.environ.get("AWS_EXECUTION_ENV") or os.environ.get("AWS_LAMBDA_FUNCTION_NAME")
+
+                client = clickhouse_connect.get_client(
+                    host=self._config.host,
+                    port=self._config.port,
+                    username=self._config.username,
+                    password=self._config.password,
+                    database=self._config.database,
+                    secure=self._config.secure,
+                    verify=not bool(is_lambda),  # Disable SSL verification in Lambda
+                )
+                # Test connection with simple query
+                client.command("SELECT 1")
+                return client
+            except Exception as e:
+                if "authentication" in str(e).lower() or "credential" in str(e).lower():
+                    # Authentication failures should not retry
+                    raise ClickHouseConnectionError(
+                        f"Authentication failed: Invalid credentials for user '{self._config.username}'. "
+                        f"Check CLICKHOUSE_USER and CLICKHOUSE_PASSWORD environment variables."
+                    ) from e
+                elif attempt < max_retries - 1:
+                    # Network errors: retry once
+                    time.sleep(retry_delay)
+                    continue
+                else:
+                    # Final attempt failed
+                    raise ClickHouseConnectionError(f"Failed to connect to ClickHouse at {self._config.host}: {e}") from e
+
+        # Should never reach here due to max_retries logic, but for type safety
+        raise ClickHouseConnectionError("Failed to establish connection after retries")
+
+    def get_config(self) -> ClickHouseConfig:
+        """Get current configuration.
+
+        Returns:
+            ClickHouseConfig instance
+
+        Raises:
+            ClickHouseConfigError: If configuration loading fails
+        """
+        if self._config is None:
+            self._config = ClickHouseConfig.load()
+        return self._config
+
+    def close(self) -> None:
+        """Close the current connection.
+
+        Note: Typically not needed for CLI use cases (process termination handles cleanup).
+        Provided for completeness and testing.
+        """
+        if self._client is not None:
+            try:
+                self._client.close()
+            except Exception:  # nosec B110 - cleanup, safe to ignore failures
+                pass  # Best effort close
+            finally:
+                self._client = None
+
+
+class ClickHouseClient:
+    """Wrapper for ClickHouse query execution with formatted output."""
+
+    def __init__(self) -> None:
+        """Initialize ClickHouse client wrapper."""
+        self.manager = ClickHouseConnectionManager.get_instance()
+
+    def execute_query(self, query: str, format: str = "json") -> Dict[str, Any]:
+        """Execute query and return formatted results.
+
+        Args:
+            query: SQL query to execute
+            format: Output format ('json', 'table', 'csv')
+
+        Returns:
+            Dictionary with query results and metadata:
+            {
+                'columns': List[str],
+                'data': List[List[Any]],
+                'rows': int,
+                'elapsed': str,
+                'query': str
+            }
+
+        Raises:
+            ClickHouseQueryError: If query execution fails
+        """
+        try:
+            client = self.manager.get_client()
+            start_time = time.time()
+
+            # Execute query
+            result = client.query(query)
+
+            elapsed = time.time() - start_time
+            elapsed_ms = int(elapsed * 1000)
+
+            # Extract column names and data
+            columns = result.column_names
+            data = result.result_rows
+            rows = len(data)
+
+            # Auto-log metrics to centralized tracker
+            try:
+                from athf.core.metrics_tracker import MetricsTracker
+
+                MetricsTracker.get_instance().log_clickhouse_query(
+                    sql=query,
+                    duration_ms=elapsed_ms,
+                    rows=rows,
+                    status="success",
+                )
+            except Exception:
+                pass  # Never fail query execution due to metrics logging
+
+            return {
+                "columns": columns,
+                "data": data,
+                "rows": rows,
+                "elapsed": f"{elapsed:.3f}s",
+                "query": query,
+            }
+
+        except Exception as e:
+            # Log error metrics
+            try:
+                from athf.core.metrics_tracker import MetricsTracker
+
+                status = "timeout" if "timeout" in str(e).lower() else "error"
+                MetricsTracker.get_instance().log_clickhouse_query(
+                    sql=query,
+                    duration_ms=0,  # Unknown duration on error
+                    rows=0,
+                    status=status,
+                )
+            except Exception:
+                pass  # Never fail due to metrics logging
+
+            # Check for timeout errors
+            if "timeout" in str(e).lower():
+                raise ClickHouseQueryError(
+                    f"Query timeout: {e}\n\n"
+                    "Tips to avoid timeouts:\n"
+                    "  1. Add time bounds: WHERE timestamp >= now() - INTERVAL 7 DAY\n"
+                    "  2. Start with small LIMIT: LIMIT 100\n"
+                    "  3. Filter early: Add WHERE clause before aggregations\n"
+                    '  4. Validate query: athf validate query --sql "..."'
+                ) from e
+            else:
+                raise ClickHouseQueryError(f"Query execution failed: {e}") from e
+
+    def test_connection(self) -> Dict[str, Any]:
+        """Test ClickHouse connection with simple query.
+
+        Returns:
+            Dictionary with connection status and details:
+            {
+                'success': bool,
+                'host': str,
+                'port': int,
+                'database': str,
+                'message': str
+            }
+
+        Raises:
+            ClickHouseConnectionError: If connection test fails
+        """
+        try:
+            client = self.manager.get_client()
+            client.command("SELECT 1")
+
+            config = self.manager.get_config()
+
+            return {
+                "success": True,
+                "host": config.host,
+                "port": config.port,
+                "database": config.database,
+                "message": "Connection successful",
+            }
+        except Exception as e:
+            config = self.manager.get_config()
+            return {
+                "success": False,
+                "host": config.host,
+                "port": config.port,
+                "database": config.database,
+                "message": f"Connection failed: {e}",
+            }