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

athf/core/query_parser.py

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

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

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