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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentic-threat-hunting-framework
-Version: 0.5.0
+Version: 0.5.1
 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.5.0.dist-info → agentic_threat_hunting_framework-0.5.1.dist-info}/RECORD

@@ -1,10 +1,10 @@
-agentic_threat_hunting_framework-0.5.0.dist-info/licenses/LICENSE,sha256=_KObErRfiKoolznt-DF0nJnr3U9Rdh7Z4Ba7G5qqckk,1071
+agentic_threat_hunting_framework-0.5.1.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=108PnDRlaytJj9KzjzcTLljB3DeerMIXZOeAQJrmtPU,5052
 athf/plugin_system.py,sha256=c_9-oUiR6tuYpWpEmeVRayU8-TXlkjvZC3EUxuYWW4M,1515
 athf/agents/__init__.py,sha256=iaSJpvnXm9rz4QS7gBrsaLEjm49uvsMs4BLPOJeyp78,346
-athf/agents/base.py,sha256=lnVDIOQUOyP-Apa9UM2E1mRXUPnNJ4hVqQXOwVw2u4c,4286
+athf/agents/base.py,sha256=ZLqSW6I0pKqs1Z3OIoV8urkysMRzNDs52yNIxDgQjTU,3981
 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
@@ -20,17 +20,10 @@ athf/commands/similar.py,sha256=FTTVr4zzP9bdJrirscp6pOxdQbE8zot6pa20-_TYiuo,1180
 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=B9IhmwH7gy2RVA6WSN3L7yGp3Q4L8OsiiwcEvnnZejU,10320
@@ -59,8 +52,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.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.5.1.dist-info/METADATA,sha256=A9gPoTr0bXjdPo_l76Fn1tIxJFAL3RFQRRUPSDT6N8I,15949
+agentic_threat_hunting_framework-0.5.1.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+agentic_threat_hunting_framework-0.5.1.dist-info/entry_points.txt,sha256=GopR2iTiBs-yNMWiUZ2DaFIFglXxWJx1XPjTa3ePtfE,39
+agentic_threat_hunting_framework-0.5.1.dist-info/top_level.txt,sha256=Cxxg6SMLfawDJWBITsciRzq27XV8fiaAor23o9Byoes,5
+agentic_threat_hunting_framework-0.5.1.dist-info/RECORD,,

athf/agents/base.py

@@ -1,4 +1,4 @@
-"""Base classes for hunt-vault agents."""
+"""Base classes for ATHF agents."""
 
 import os
 from abc import ABC, abstractmethod
@@ -89,6 +89,8 @@ class LLMAgent(Agent[InputT, OutputT]):
     ) -> None:
         """Log LLM call metrics to centralized tracker.
 
+        Override this method in subclasses or plugins to implement custom metrics tracking.
+
         Args:
             agent_name: Name of the agent (e.g., "hypothesis-generator")
             model_id: Bedrock model ID
@@ -97,19 +99,8 @@ class LLMAgent(Agent[InputT, OutputT]):
             cost_usd: Estimated cost in USD
             duration_ms: Call duration in milliseconds
         """
-        try:
-            from athf.core.metrics_tracker import MetricsTracker
-
-            MetricsTracker.get_instance().log_bedrock_call(
-                agent=agent_name,
-                model_id=model_id,
-                input_tokens=input_tokens,
-                output_tokens=output_tokens,
-                cost_usd=cost_usd,
-                duration_ms=duration_ms,
-            )
-        except Exception:
-            pass  # Never fail agent execution due to metrics logging
+        # No-op by default. Override in plugins for custom metrics tracking.
+        pass
 
     def _get_llm_client(self) -> Any:
         """Get AWS Bedrock runtime client for Claude models.

athf/core/clickhouse_connection.py

@@ -1,396 +0,0 @@
-"""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}",
-            }