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

athf/core/query_parser.py

@@ -1,203 +0,0 @@
-"""Query library parser for YAML query files."""
-
-import re
-from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
-
-import yaml
-
-
-class QueryParser:
-    """Parser for query library YAML files."""
-
-    def __init__(self, queries_dir: Path):
-        """Initialize query parser with queries directory.
-
-        Args:
-            queries_dir: Path to queries/ directory containing YAML files
-        """
-        self.queries_dir = Path(queries_dir)
-        self._queries_cache: Dict[str, Dict[str, Any]] = {}
-        self._loaded = False
-
-    def load_all_queries(self) -> List[Dict[str, Any]]:
-        """Load all queries from YAML files in queries/ directory.
-
-        Returns:
-            List of query dictionaries
-        """
-        if self._loaded and self._queries_cache:
-            return list(self._queries_cache.values())
-
-        all_queries = []
-        yaml_files = sorted(self.queries_dir.glob("*.yaml"))
-
-        for yaml_file in yaml_files:
-            try:
-                with open(yaml_file, "r", encoding="utf-8") as f:
-                    data = yaml.safe_load(f)
-
-                if not data or not isinstance(data, dict):
-                    continue
-
-                category = data.get("category")
-                queries = data.get("queries", [])
-
-                for query in queries:
-                    # Add category from file-level metadata
-                    query["category"] = category
-                    query["source_file"] = yaml_file.name
-
-                    # Cache by query_id
-                    query_id = query.get("query_id")
-                    if query_id:
-                        self._queries_cache[query_id] = query
-
-                    all_queries.append(query)
-
-            except Exception as e:
-                # Log error but continue loading other files
-                print(f"Warning: Failed to parse {yaml_file}: {e}")
-                continue
-
-        self._loaded = True
-        return all_queries
-
-    def get_query_by_id(self, query_id: str) -> Optional[Dict[str, Any]]:
-        """Retrieve query metadata by Q-XXXX ID.
-
-        Args:
-            query_id: Query ID (e.g., Q-USER-001)
-
-        Returns:
-            Query dictionary or None if not found
-        """
-        if not self._loaded:
-            self.load_all_queries()
-
-        return self._queries_cache.get(query_id)
-
-    def search_queries(self, keyword: str) -> List[Dict[str, Any]]:
-        """Search queries by keyword in name, description, tags.
-
-        Args:
-            keyword: Search keyword
-
-        Returns:
-            List of matching query dictionaries
-        """
-        all_queries = self.load_all_queries()
-        keyword_lower = keyword.lower()
-        results = []
-
-        for query in all_queries:
-            # Search in name
-            if keyword_lower in query.get("name", "").lower():
-                results.append(query)
-                continue
-
-            # Search in description
-            if keyword_lower in query.get("description", "").lower():
-                results.append(query)
-                continue
-
-            # Search in tags
-            tags = query.get("tags", [])
-            if any(keyword_lower in tag.lower() for tag in tags):
-                results.append(query)
-                continue
-
-        return results
-
-    def filter_queries(self, category: Optional[str] = None, tags: Optional[List[str]] = None) -> List[Dict[str, Any]]:
-        """Filter queries by category or tags.
-
-        Args:
-            category: Category to filter by (e.g., "user-activity")
-            tags: List of tags to filter by (any match)
-
-        Returns:
-            List of matching query dictionaries
-        """
-        all_queries = self.load_all_queries()
-        results = all_queries
-
-        if category:
-            results = [q for q in results if q.get("category") == category]
-
-        if tags:
-            tag_set = {t.lower() for t in tags}
-            results = [q for q in results if any(qt.lower() in tag_set for qt in q.get("tags", []))]
-
-        return results
-
-    def validate_query(self, query: Dict[str, Any]) -> Tuple[bool, List[str]]:
-        """Validate query structure and required fields.
-
-        Args:
-            query: Query dictionary
-
-        Returns:
-            Tuple of (is_valid, list_of_errors)
-        """
-        errors = []
-
-        # Required fields
-        required_fields = ["query_id", "name", "description", "query"]
-        for field in required_fields:
-            if not query.get(field):
-                errors.append(f"Missing required field: {field}")
-
-        # Validate query_id format: Q-[A-Z]+-[0-9]+
-        query_id = query.get("query_id", "")
-        if not re.match(r"^Q-[A-Z]+-\d+$", query_id):
-            errors.append(f"Invalid query_id format: {query_id} (expected Q-[CATEGORY]-[NUMBER])")
-
-        # Validate query includes time bounds
-        query_sql = query.get("query", "")
-        has_time_bound = "INTERVAL" in query_sql.upper() or "timestamp >=" in query_sql or "timestamp BETWEEN" in query_sql
-        if not has_time_bound:
-            errors.append("Query missing time constraint (INTERVAL or timestamp >=)")
-
-        # Validate query includes LIMIT clause
-        if "LIMIT" not in query_sql.upper():
-            errors.append("Query missing LIMIT clause")
-
-        # Validate placeholders match {{}} in query
-        placeholders_defined = set(query.get("placeholders", {}).keys())
-        placeholders_in_query = set(re.findall(r"\{\{(\w+)\}\}", query_sql))
-
-        # Check for placeholders used in query but not defined
-        undefined = placeholders_in_query - placeholders_defined
-        if undefined:
-            errors.append(f"Placeholders used but not defined: {', '.join(undefined)}")
-
-        # Check for placeholders defined but not used (warning, not error)
-        unused = placeholders_defined - placeholders_in_query
-        if unused:
-            # This is just a warning, not an error
-            pass
-
-        return (len(errors) == 0, errors)
-
-    def get_categories(self) -> List[str]:
-        """Get list of all unique categories in query library.
-
-        Returns:
-            List of category names
-        """
-        all_queries = self.load_all_queries()
-        categories: set[str] = {cat for q in all_queries if (cat := q.get("category")) is not None}
-        return sorted(categories)
-
-    def get_all_tags(self) -> List[str]:
-        """Get list of all unique tags in query library.
-
-        Returns:
-            List of tag names
-        """
-        all_queries = self.load_all_queries()
-        tags = set()
-        for query in all_queries:
-            tags.update(query.get("tags", []))
-        return sorted(tags)

athf/core/query_suggester.py

@@ -1,235 +0,0 @@
-"""Query suggestion engine for LLM-driven alert triage."""
-
-import re
-from pathlib import Path
-from typing import Any, Dict, List, Tuple
-
-from athf.core.query_parser import QueryParser
-
-
-class QuerySuggester:
-    """Suggests relevant queries based on alert text analysis."""
-
-    def __init__(self, queries_dir: Path):
-        """Initialize query suggester.
-
-        Args:
-            queries_dir: Path to queries/ directory
-        """
-        self.parser = QueryParser(queries_dir)
-        self.all_queries = self.parser.load_all_queries()
-
-    def extract_parameters(self, alert_text: str) -> Dict[str, str]:
-        """Extract username, hostname, and other parameters from alert text.
-
-        Patterns supported:
-        - "user: john.doe" or "username: john.doe" or "User: john.doe"
-        - "host: LAPTOP-123" or "hostname: LAPTOP-123" or "Host: LAPTOP-123"
-        - "process: powershell.exe" or "Process: powershell.exe"
-        - "ip: 8.8.8.8" or "IP: 8.8.8.8"
-        - "organization: acme-corp" or "org: acme-corp"
-
-        Args:
-            alert_text: Alert text to analyze
-
-        Returns:
-            Dictionary of extracted parameters
-        """
-        parameters = {}
-
-        # Username patterns
-        username_patterns = [
-            r"(?:user|username|User|USERNAME):\s*([a-zA-Z0-9._@-]+)",
-            r"user\s+([a-zA-Z0-9._@-]+)",
-            r"for\s+user\s+([a-zA-Z0-9._@-]+)",
-        ]
-        for pattern in username_patterns:
-            if match := re.search(pattern, alert_text):
-                parameters["username"] = match.group(1)
-                break
-
-        # Hostname patterns
-        hostname_patterns = [
-            r"(?:host|hostname|Host|HOSTNAME):\s*([a-zA-Z0-9._-]+)",
-            r"(?:on|from)\s+host\s+([a-zA-Z0-9._-]+)",
-            r"endpoint\s+([a-zA-Z0-9._-]+)",
-        ]
-        for pattern in hostname_patterns:
-            if match := re.search(pattern, alert_text):
-                parameters["hostname"] = match.group(1)
-                break
-
-        # Process name patterns
-        process_patterns = [
-            r"(?:process|Process|PROCESS):\s*([a-zA-Z0-9._-]+\.exe)",
-            r"(?:process|Process)\s+([a-zA-Z0-9._-]+\.exe)",
-            r"execution\s+of\s+([a-zA-Z0-9._-]+\.exe)",
-        ]
-        for pattern in process_patterns:
-            if match := re.search(pattern, alert_text):
-                parameters["process_name"] = match.group(1)
-                break
-
-        # IP address patterns
-        ip_pattern = r"(?:ip|IP|address):\s*(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})"
-        if match := re.search(ip_pattern, alert_text):
-            parameters["ip_address"] = match.group(1)
-
-        # Organization patterns
-        org_patterns = [
-            r"(?:organization|org):\s*([a-zA-Z0-9_-]+)",
-            r"customer\s+([a-zA-Z0-9_-]+)",
-        ]
-        for pattern in org_patterns:
-            if match := re.search(pattern, alert_text):
-                parameters["organization_id"] = match.group(1)
-                break
-
-        return parameters
-
-    def suggest_queries(self, alert_text: str, max_suggestions: int = 5) -> List[Tuple[int, Dict[str, Any]]]:
-        """Analyze alert text and suggest relevant queries with scores.
-
-        Scoring algorithm:
-        - Keyword in query name: +3 points
-        - Keyword in query description: +2 points
-        - Keyword in query tags: +1 point
-        - Extracted parameter matches query placeholder: +5 points
-
-        Args:
-            alert_text: Alert text to analyze
-            max_suggestions: Maximum number of queries to suggest
-
-        Returns:
-            List of tuples (score, query) sorted by relevance
-        """
-        # Extract parameters
-        parameters = self.extract_parameters(alert_text)
-
-        # Extract keywords
-        keywords = self._extract_keywords(alert_text)
-
-        # Score queries by relevance
-        scored_queries = []
-        for query in self.all_queries:
-            score = 0
-
-            # Match keywords against query metadata
-            for keyword in keywords:
-                keyword_lower = keyword.lower()
-
-                # Check name
-                if keyword_lower in query.get("name", "").lower():
-                    score += 3
-
-                # Check description
-                if keyword_lower in query.get("description", "").lower():
-                    score += 2
-
-                # Check tags
-                query_tags = " ".join(query.get("tags", [])).lower()
-                if keyword_lower in query_tags:
-                    score += 1
-
-            # Boost score if extracted parameters match query placeholders
-            query_placeholders = set(query.get("placeholders", {}).keys())
-            for param in parameters.keys():
-                if param in query_placeholders:
-                    score += 5
-
-            if score > 0:
-                scored_queries.append((score, query))
-
-        # Sort by score (descending) and return top N
-        scored_queries.sort(key=lambda x: x[0], reverse=True)
-        return scored_queries[:max_suggestions]
-
-    def _extract_keywords(self, text: str) -> List[str]:
-        """Extract relevant security keywords from alert text.
-
-        Args:
-            text: Alert text to analyze
-
-        Returns:
-            List of security-related keywords found
-        """
-        keywords = []
-        text_lower = text.lower()
-
-        # Security keywords organized by category
-        security_terms = {
-            # Execution
-            "powershell",
-            "cmd",
-            "bash",
-            "shell",
-            "script",
-            "execution",
-            "process",
-            # Credential Access
-            "credential",
-            "password",
-            "mimikatz",
-            "lsass",
-            "procdump",
-            "sam",
-            "ntds",
-            # Network
-            "network",
-            "connection",
-            "lateral",
-            "smb",
-            "rdp",
-            "winrm",
-            "exfiltration",
-            "c2",
-            "command-and-control",
-            # General
-            "suspicious",
-            "malicious",
-            "privilege",
-            "escalation",
-            "file",
-            "access",
-        }
-
-        for term in security_terms:
-            if term in text_lower:
-                keywords.append(term)
-
-        # Extract MITRE ATT&CK technique IDs
-        technique_pattern = r"T\d{4}(?:\.\d{3})?"
-        techniques = re.findall(technique_pattern, text, re.IGNORECASE)
-        keywords.extend(techniques)
-
-        return keywords
-
-    def get_parameter_coverage(self, query: Dict[str, Any], extracted_params: Dict[str, str]) -> float:
-        """Calculate what percentage of query placeholders can be filled.
-
-        Args:
-            query: Query dictionary
-            extracted_params: Extracted parameters from alert
-
-        Returns:
-            Coverage percentage (0.0 to 1.0)
-        """
-        placeholders = query.get("placeholders", {})
-        if not placeholders:
-            return 1.0  # No placeholders, fully covered
-
-        # Count how many placeholders can be filled
-        required = [name for name, info in placeholders.items() if "default" not in info]
-        optional = [name for name, info in placeholders.items() if "default" in info]
-
-        filled_required = sum(1 for r in required if r in extracted_params)
-        filled_optional = sum(1 for o in optional if o in extracted_params)
-
-        # Weight required higher than optional
-        total_score = len(required) * 2 + len(optional)
-        filled_score = filled_required * 2 + filled_optional
-
-        if total_score == 0:
-            return 1.0
-
-        return filled_score / total_score

athf/core/query_validator.py

@@ -1,240 +0,0 @@
-"""Query validation for ClickHouse SQL safety checks."""
-
-import re
-from dataclasses import dataclass
-from typing import List, Optional
-
-
-@dataclass
-class ValidationResult:
-    """Result of query validation."""
-
-    is_valid: bool
-    errors: List[str]
-    warnings: List[str]
-    info: List[str]
-    estimated_row_count: Optional[int] = None
-
-
-class QueryValidator:
-    """Validates SQL queries for ClickHouse MCP safety requirements.
-
-    Enforces rules from integrations/clickhouse/README.md:
-    - Time bounds required (7 days max initially)
-    - LIMIT clause required (start with 100)
-    - No SQL comments (breaks MCP)
-    - Early filtering recommended
-    - Avoid multiple DISTINCT
-    """
-
-    # Time-related patterns
-    TIME_PATTERNS = [
-        r"toDate\(.+\)\s*>=\s*",  # toDate(timestamp) >= ...
-        r"toDateTime\(.+\)\s*>=\s*",  # toDateTime(timestamp) >= ...
-        r"timestamp\s*>=\s*",  # timestamp >= ...
-        r"time\s*>=\s*",  # time >= ...
-        r"WHERE.+>=\s*now\(\)\s*-\s*INTERVAL",  # WHERE ... >= now() - INTERVAL
-        r"WHERE.+>=\s*subtractDays",  # WHERE ... >= subtractDays()
-        r"WHERE.+>=\s*subtractHours",  # WHERE ... >= subtractHours()
-        r">=\s*toDateTime\(",  # >= toDateTime('...')
-        r">=\s*'[\d\-:\s]+'",  # >= '2025-01-15 00:00:00' or >= '2025-01-15'
-        r"<=\s*'[\d\-:\s]+'",  # <= '2025-01-15 23:59:59'
-        r"BETWEEN\s+['\"]\d{4}-\d{2}-\d{2}",  # BETWEEN '2025-01-15' AND '2025-01-22'
-        r"=\s*'[\d\-:\s]+'",  # = '2025-01-15' (exact date match)
-    ]
-
-    # LIMIT patterns
-    LIMIT_PATTERNS = [
-        r"LIMIT\s+\d+",  # LIMIT 100
-        r"TOP\s+\d+",  # TOP 100 (alternative syntax)
-    ]
-
-    # Comment patterns (these break MCP)
-    COMMENT_PATTERNS = [
-        r"--",  # Single-line comment
-        r"/\*.*?\*/",  # Multi-line comment
-    ]
-
-    # Expensive operations
-    EXPENSIVE_PATTERNS = [
-        (r"DISTINCT.*DISTINCT", "Multiple DISTINCT clauses detected (expensive operation)"),
-        (r"SELECT\s+\*\s+FROM", "SELECT * detected (prefer explicit columns for performance)"),
-        (r"GROUP BY.+ORDER BY.+LIMIT", "GROUP BY + ORDER BY without early filtering may be slow"),
-    ]
-
-    def __init__(self) -> None:
-        """Initialize the query validator."""
-        pass
-
-    def validate(self, query: str, target: str = "clickhouse") -> ValidationResult:
-        """Validate a SQL query against safety requirements.
-
-        Args:
-            query: The SQL query to validate
-            target: The target database (currently only 'clickhouse' supported)
-
-        Returns:
-            ValidationResult with validation status and messages
-        """
-        if target != "clickhouse":
-            return ValidationResult(
-                is_valid=False,
-                errors=[f"Unsupported target: {target} (only 'clickhouse' is currently supported)"],
-                warnings=[],
-                info=[],
-            )
-
-        errors: List[str] = []
-        warnings: List[str] = []
-        info: List[str] = []
-
-        # Normalize query (remove extra whitespace)
-        normalized = " ".join(query.split())
-
-        # 1. CRITICAL: Check for SQL comments (breaks MCP)
-        if self._has_comments(query):
-            errors.append("SQL comments detected (-- or /* */). Comments break MCP - remove them before execution.")
-
-        # 2. CRITICAL: Check for time bounds
-        if not self._has_time_bounds(normalized):
-            errors.append(
-                "No time bounds detected. Add time constraints (e.g., WHERE timestamp >= now() - INTERVAL 7 DAY "
-                "or WHERE time >= '2025-01-15 00:00:00') to prevent timeouts."
-            )
-        else:
-            info.append("Time bounds detected")
-
-        # 3. CRITICAL: Check for LIMIT clause
-        limit_value = self._extract_limit(normalized)
-        if limit_value is None:
-            errors.append("No LIMIT clause detected. Add LIMIT (start with 100, max 1000) to prevent large result sets.")
-        else:
-            if limit_value <= 100:
-                info.append(f"LIMIT {limit_value} (safe - good starting point)")
-            elif limit_value <= 1000:
-                info.append(f"LIMIT {limit_value} (moderate - may need refinement if exactly {limit_value} rows returned)")
-            else:
-                warnings.append(
-                    f"LIMIT {limit_value} is high. Consider starting with LIMIT 100 and increasing if needed "
-                    "(progressive strategy)."
-                )
-
-        # 4. Check for early filtering (WHERE clause before aggregations)
-        if not self._has_early_filtering(normalized):
-            warnings.append(
-                "No early filtering detected. Add WHERE clause before aggregations (GROUP BY, DISTINCT) "
-                "for better performance."
-            )
-        else:
-            info.append("Early filtering detected (WHERE clause present)")
-
-        # 5. Check for expensive operations
-        for pattern, message in self.EXPENSIVE_PATTERNS:
-            if re.search(pattern, normalized, re.IGNORECASE):
-                warnings.append(message)
-
-        # 6. Check for query structure (FROM clause required)
-        if not re.search(r"FROM\s+\w+", normalized, re.IGNORECASE):
-            errors.append("No FROM clause detected. Query must select from a table.")
-
-        # Determine overall validity
-        is_valid = len(errors) == 0
-
-        return ValidationResult(
-            is_valid=is_valid,
-            errors=errors,
-            warnings=warnings,
-            info=info,
-        )
-
-    def _has_comments(self, query: str) -> bool:
-        """Check if query contains SQL comments."""
-        # Check for single-line comments
-        if "--" in query:
-            return True
-
-        # Check for multi-line comments
-        if re.search(r"/\*.*?\*/", query, re.DOTALL):
-            return True
-
-        return False
-
-    def _has_time_bounds(self, normalized_query: str) -> bool:
-        """Check if query has time constraints."""
-        for pattern in self.TIME_PATTERNS:
-            if re.search(pattern, normalized_query, re.IGNORECASE):
-                return True
-        return False
-
-    def _extract_limit(self, normalized_query: str) -> Optional[int]:
-        """Extract LIMIT value from query."""
-        # Check for LIMIT clause
-        limit_match = re.search(r"LIMIT\s+(\d+)", normalized_query, re.IGNORECASE)
-        if limit_match:
-            return int(limit_match.group(1))
-
-        # Check for TOP clause (alternative syntax)
-        top_match = re.search(r"TOP\s+(\d+)", normalized_query, re.IGNORECASE)
-        if top_match:
-            return int(top_match.group(1))
-
-        return None
-
-    def _has_early_filtering(self, normalized_query: str) -> bool:
-        """Check if query has WHERE clause before aggregations."""
-        # Look for WHERE clause
-        where_match = re.search(r"WHERE", normalized_query, re.IGNORECASE)
-        if not where_match:
-            return False
-
-        # Check if WHERE comes before GROUP BY or DISTINCT
-        where_pos = where_match.start()
-
-        group_match = re.search(r"GROUP\s+BY", normalized_query, re.IGNORECASE)
-        if group_match and group_match.start() < where_pos:
-            return False
-
-        distinct_match = re.search(r"DISTINCT", normalized_query, re.IGNORECASE)
-        if distinct_match and distinct_match.start() < where_pos:
-            return False
-
-        return True
-
-    def suggest_improvements(self, query: str, validation: ValidationResult) -> List[str]:
-        """Suggest improvements based on validation results.
-
-        Args:
-            query: The original query
-            validation: The validation result
-
-        Returns:
-            List of suggested improvements
-        """
-        suggestions: List[str] = []
-
-        # If missing time bounds, suggest adding them
-        if any("time bounds" in error.lower() for error in validation.errors):
-            suggestions.append(
-                "Add time bounds: WHERE timestamp >= now() - INTERVAL 7 DAY "
-                "or WHERE time >= '2025-01-15 00:00:00' AND time <= '2025-01-22 23:59:59'"
-            )
-
-        # If missing LIMIT, suggest adding it
-        if any("limit" in error.lower() for error in validation.errors):
-            suggestions.append("Add LIMIT clause: LIMIT 100 (progressive strategy)")
-
-        # If has comments, suggest removing them
-        if any("comment" in error.lower() for error in validation.errors):
-            suggestions.append("Remove SQL comments (-- or /* */) - they break ClickHouse MCP")
-
-        # If missing early filtering
-        if any("early filtering" in warning.lower() for warning in validation.warnings):
-            suggestions.append(
-                "Add WHERE clause before aggregations for better performance: " "WHERE timestamp >= ... AND condition"
-            )
-
-        # If SELECT *
-        if any("SELECT \\*" in warning for warning in validation.warnings):
-            suggestions.append("Replace SELECT * with explicit columns: SELECT column1, column2, ...")
-
-        return suggestions