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

athf/commands/splunk.py

@@ -0,0 +1,323 @@
+"""Splunk integration commands for ATHF.
+
+This module provides CLI commands for interacting with Splunk via REST API.
+"""
+
+import json
+import os
+from typing import Optional
+
+import click
+from rich.console import Console
+from rich.table import Table
+
+from athf.core.splunk_client import SplunkClient
+
+console = Console()
+
+
+def get_client(host: Optional[str], token: Optional[str], verify_ssl: Optional[bool]) -> SplunkClient:
+    """Get Splunk client from CLI args or environment variables.
+
+    Args:
+        host: Splunk host (from CLI)
+        token: Auth token (from CLI)
+        verify_ssl: Whether to verify SSL (None to read from env)
+
+    Returns:
+        Configured SplunkClient
+
+    Raises:
+        click.UsageError: If credentials are not provided
+    """
+    # Try CLI args first, fall back to environment
+    if not host:
+        host = os.getenv("SPLUNK_HOST")
+    if not token:
+        token = os.getenv("SPLUNK_TOKEN")
+
+    # Read verify_ssl from environment if not provided via CLI
+    if verify_ssl is None:
+        env_verify = os.getenv("SPLUNK_VERIFY_SSL", "true")
+        verify_ssl = env_verify.lower() in ("true", "1", "yes")
+
+    if not host or not token:
+        raise click.UsageError(
+            "Splunk credentials required. Provide via:\n"
+            "  • CLI: --host and --token flags\n"
+            "  • Environment: SPLUNK_HOST and SPLUNK_TOKEN variables\n"
+            "  • Config file: Create .env with credentials"
+        )
+
+    return SplunkClient(host=host, token=token, verify_ssl=verify_ssl)
+
+
+@click.group()
+def splunk() -> None:
+    """Splunk REST API integration.
+
+    \b
+    Execute SPL queries and interact with Splunk directly from ATHF CLI.
+
+    \b
+    Setup:
+      1. Create a Splunk authentication token:
+         Settings → Tokens → New Token
+
+      2. Set environment variables (recommended):
+         export SPLUNK_HOST="splunk.example.com"
+         export SPLUNK_TOKEN="your-token-here"
+
+      3. Or use --host and --token flags with each command
+
+    \b
+    Examples:
+      # Test connection
+      athf splunk test
+
+      # List available indexes
+      athf splunk indexes
+
+      # Execute a query
+      athf splunk search 'index=main "Failed password" | head 10'
+
+      # Query with time range
+      athf splunk search 'index=* | stats count by sourcetype' \\
+        --earliest "-7d" --latest "now" --count 100
+    """
+
+
+@splunk.command()
+@click.option("--host", envvar="SPLUNK_HOST", help="Splunk host (e.g., splunk.example.com)")
+@click.option("--token", envvar="SPLUNK_TOKEN", help="Splunk authentication token")
+@click.option("--verify-ssl/--no-verify-ssl", default=None, help="Verify SSL certificates")
+def test(host: Optional[str], token: Optional[str], verify_ssl: Optional[bool]) -> None:
+    """Test Splunk connection and authentication.
+
+    \b
+    Validates that:
+      • Host is reachable
+      • Token is valid
+      • API access is working
+
+    \b
+    Example:
+      athf splunk test
+    """
+    try:
+        client = get_client(host, token, verify_ssl)
+        info = client.test_connection()
+
+        console.print("\n[bold green]✓ Connection successful![/bold green]\n")
+
+        # Display server info
+        if "entry" in info and len(info["entry"]) > 0:
+            content = info["entry"][0].get("content", {})
+            console.print(f"[bold]Server:[/bold] {content.get('serverName', 'N/A')}")
+            console.print(f"[bold]Version:[/bold] {content.get('version', 'N/A')}")
+            console.print(f"[bold]Build:[/bold] {content.get('build', 'N/A')}")
+
+    except Exception as e:
+        console.print(f"\n[bold red]✗ Connection failed:[/bold red] {e}\n", style="red")
+        raise click.Abort()
+
+
+@splunk.command()
+@click.option("--host", envvar="SPLUNK_HOST", help="Splunk host")
+@click.option("--token", envvar="SPLUNK_TOKEN", help="Splunk authentication token")
+@click.option("--verify-ssl/--no-verify-ssl", default=None, help="Verify SSL certificates")
+@click.option("--format", "output_format", type=click.Choice(["table", "json", "list"]), default="list", help="Output format")
+def indexes(host: Optional[str], token: Optional[str], verify_ssl: Optional[bool], output_format: str) -> None:
+    """List available Splunk indexes.
+
+    \b
+    Shows all indexes accessible with current credentials.
+
+    \b
+    Example:
+      athf splunk indexes
+      athf splunk indexes --format json
+    """
+    try:
+        client = get_client(host, token, verify_ssl)
+        index_list = client.get_indexes()
+
+        if not index_list:
+            console.print("[yellow]No indexes found[/yellow]")
+            return
+
+        if output_format == "json":
+            click.echo(json.dumps({"indexes": index_list}, indent=2))
+        elif output_format == "table":
+            table = Table(title=f"Splunk Indexes ({len(index_list)} total)")
+            table.add_column("Index Name", style="cyan")
+            for idx in sorted(index_list):
+                table.add_row(idx)
+            console.print(table)
+        else:  # list
+            console.print(f"\n[bold]Available Indexes ({len(index_list)}):[/bold]\n")
+            for idx in sorted(index_list):
+                console.print(f"  • {idx}")
+            console.print()
+
+    except Exception as e:
+        console.print(f"[bold red]Error:[/bold red] {e}", style="red")
+        raise click.Abort()
+
+
+@splunk.command()
+@click.argument("query")
+@click.option("--host", envvar="SPLUNK_HOST", help="Splunk host")
+@click.option("--token", envvar="SPLUNK_TOKEN", help="Splunk authentication token")
+@click.option("--verify-ssl/--no-verify-ssl", default=None, help="Verify SSL certificates")
+@click.option("--earliest", default="-24h", help="Earliest time (e.g., '-24h', '2024-01-01T00:00:00')")
+@click.option("--latest", default="now", help="Latest time (e.g., 'now', '2024-01-02T00:00:00')")
+@click.option("--count", default=100, type=int, help="Maximum results to return")
+@click.option("--format", "output_format", type=click.Choice(["json", "table", "raw"]), default="json", help="Output format")
+@click.option("--async-search/--oneshot", "use_async", default=False, help="Use async search for long queries")
+@click.option("--max-wait", default=300, type=int, help="Max wait time for async searches (seconds)")
+def search(
+    query: str,
+    host: Optional[str],
+    token: Optional[str],
+    verify_ssl: Optional[bool],
+    earliest: str,
+    latest: str,
+    count: int,
+    output_format: str,
+    use_async: bool,
+    max_wait: int,
+) -> None:
+    """Execute a Splunk search query.
+
+    \b
+    Runs SPL (Splunk Processing Language) queries and returns results.
+
+    \b
+    Query Examples:
+      'index=main "Failed password"'
+      'index=* sourcetype=linux_secure | stats count by user'
+      'index=web status>=400 | timechart count by status'
+
+    \b
+    Time Format Examples:
+      --earliest "-1h"          (last hour)
+      --earliest "-7d"          (last 7 days)
+      --earliest "2024-01-01T00:00:00"  (absolute time)
+
+    \b
+    Examples:
+      # Basic search
+      athf splunk search 'index=main error'
+
+      # With time range
+      athf splunk search 'index=* | stats count by sourcetype' \\
+        --earliest "-7d" --count 1000
+
+      # JSON output for parsing
+      athf splunk search 'index=main | head 10' --format json
+
+      # Long-running query (async)
+      athf splunk search 'index=* | rare limit=20 sourcetype' \\
+        --async-search --max-wait 600
+    """
+    try:
+        client = get_client(host, token, verify_ssl)
+
+        console.print(f"\n[bold]Executing query:[/bold] {query}")
+        console.print(f"[bold]Time range:[/bold] {earliest} to {latest}")
+        console.print(f"[bold]Max results:[/bold] {count}\n")
+
+        # Execute search
+        if use_async:
+            console.print("[dim]Using async search (for longer queries)...[/dim]")
+            results = client.search_async(
+                query=query, earliest_time=earliest, latest_time=latest, max_results=count, max_wait=max_wait
+            )
+        else:
+            console.print("[dim]Using oneshot search (fast for small queries)...[/dim]")
+            results = client.search(query=query, earliest_time=earliest, latest_time=latest, max_count=count)
+
+        if not results:
+            console.print("[yellow]No results found[/yellow]")
+            return
+
+        console.print(f"[green]✓ Found {len(results)} results[/green]\n")
+
+        # Output results
+        if output_format == "json":
+            click.echo(json.dumps(results, indent=2, default=str))
+        elif output_format == "table":
+            if not results:
+                return
+
+            # Extract all unique fields
+            all_fields: set[str] = set()
+            for result in results:
+                all_fields.update(result.keys())
+
+            # Create table
+            table = Table(title=f"Search Results ({len(results)} events)")
+            for field in sorted(all_fields):
+                table.add_column(field, overflow="fold")
+
+            # Add rows
+            for result in results[:count]:  # Limit display
+                row = [str(result.get(field, "")) for field in sorted(all_fields)]
+                table.add_row(*row)
+
+            console.print(table)
+        else:  # raw
+            for i, result in enumerate(results, 1):
+                console.print(f"[bold cyan]Event {i}:[/bold cyan]")
+                for key, value in result.items():
+                    console.print(f"  {key}: {value}")
+                console.print()
+
+    except TimeoutError as e:
+        console.print(f"\n[bold red]Timeout:[/bold red] {e}", style="red")
+        console.print("[yellow]Try using --async-search for long queries[/yellow]")
+        raise click.Abort()
+    except Exception as e:
+        console.print(f"\n[bold red]Error:[/bold red] {e}", style="red")
+        raise click.Abort()
+
+
+@splunk.command()
+@click.option("--host", envvar="SPLUNK_HOST", help="Splunk host")
+@click.option("--token", envvar="SPLUNK_TOKEN", help="Splunk authentication token")
+@click.option("--verify-ssl/--no-verify-ssl", default=None, help="Verify SSL certificates")
+def config(host: Optional[str], token: Optional[str], verify_ssl: Optional[bool]) -> None:
+    """Show current Splunk configuration.
+
+    \b
+    Displays configuration from environment variables and validates credentials.
+
+    \b
+    Example:
+      athf splunk config
+    """
+    console.print("\n[bold]Splunk Configuration:[/bold]\n")
+
+    # Check environment
+    env_host = os.getenv("SPLUNK_HOST")
+    env_token = os.getenv("SPLUNK_TOKEN")
+    env_verify = os.getenv("SPLUNK_VERIFY_SSL", "true")
+
+    console.print(f"[bold]SPLUNK_HOST:[/bold] {env_host or '[red]Not set[/red]'}")
+    console.print(f"[bold]SPLUNK_TOKEN:[/bold] {'[green]Set[/green]' if env_token else '[red]Not set[/red]'}")
+    console.print(f"[bold]SPLUNK_VERIFY_SSL:[/bold] {env_verify}")
+
+    # Test connection if credentials available
+    if (host or env_host) and (token or env_token):
+        console.print("\n[dim]Testing connection...[/dim]")
+        try:
+            # get_client will read environment variable if verify_ssl is None
+            client = get_client(host, token, verify_ssl)
+            client.test_connection()
+            console.print("[bold green]✓ Connection successful[/bold green]\n")
+        except Exception as e:
+            console.print(f"[bold red]✗ Connection failed:[/bold red] {e}\n")
+    else:
+        console.print("\n[yellow]⚠ Missing credentials - cannot test connection[/yellow]\n")
+        console.print("[dim]Set SPLUNK_HOST and SPLUNK_TOKEN environment variables[/dim]\n")

athf/core/splunk_client.py

@@ -0,0 +1,360 @@
+"""Splunk REST API client for ATHF.
+
+This module provides direct Splunk API integration using authentication tokens.
+Use this when MCP integration is not available or for programmatic access.
+"""
+
+import time
+from typing import Any, Dict, List, Optional
+from urllib.parse import urljoin
+
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+
+class SplunkClient:
+    """Client for Splunk REST API operations.
+
+    Args:
+        host: Splunk host (e.g., "splunk.example.com" or "https://splunk.example.com:8089")
+        token: Splunk authentication token
+        verify_ssl: Whether to verify SSL certificates (default: True)
+        timeout: Request timeout in seconds (default: 30)
+
+    Example:
+        >>> client = SplunkClient(host="splunk.example.com", token="your-token")
+        >>> results = client.search("index=main | head 10", max_count=10)
+        >>> for event in results:
+        ...     print(event)
+    """
+
+    def __init__(self, host: str, token: str, verify_ssl: bool = True, timeout: int = 30):
+        # Normalize host URL
+        if not host.startswith(("http://", "https://")):
+            host = f"https://{host}"
+        if ":8089" not in host and not host.endswith(":8089"):
+            # Add default management port if not specified
+            host = host.rstrip("/") + ":8089"
+
+        self.base_url = host.rstrip("/")
+        self.token = token
+        self.verify_ssl = verify_ssl
+        self.timeout = timeout
+
+        # Create session with retry logic
+        self.session = requests.Session()
+        retry_strategy = Retry(
+            total=3,
+            backoff_factor=1,
+            status_forcelist=[429, 500, 502, 503, 504],
+        )
+        adapter = HTTPAdapter(max_retries=retry_strategy)
+        self.session.mount("http://", adapter)
+        self.session.mount("https://", adapter)
+
+        # Set default headers
+        self.session.headers.update(
+            {
+                "Authorization": f"Bearer {token}",
+                "Content-Type": "application/x-www-form-urlencoded",
+            }
+        )
+
+    def _request(
+        self,
+        method: str,
+        endpoint: str,
+        params: Optional[Dict[str, Any]] = None,
+        data: Optional[Dict[str, Any]] = None,
+        json_response: bool = True,
+    ) -> Any:
+        """Make HTTP request to Splunk API.
+
+        Args:
+            method: HTTP method (GET, POST, DELETE)
+            endpoint: API endpoint path
+            params: Query parameters
+            data: Form data for POST requests
+            json_response: Whether to parse JSON response
+
+        Returns:
+            Response data (parsed JSON or raw response)
+
+        Raises:
+            requests.HTTPError: If request fails
+        """
+        url = urljoin(self.base_url, endpoint)
+
+        response = self.session.request(
+            method=method, url=url, params=params, data=data, verify=self.verify_ssl, timeout=self.timeout
+        )
+
+        response.raise_for_status()
+
+        if json_response:
+            return response.json()
+        return response
+
+    def test_connection(self) -> Dict[str, Any]:
+        """Test connection and authentication to Splunk.
+
+        Returns:
+            Dict with server info if successful
+
+        Raises:
+            requests.HTTPError: If authentication fails
+        """
+        return self._request("GET", "/services/server/info", params={"output_mode": "json"})  # type: ignore[no-any-return]
+
+    def get_indexes(self) -> List[str]:
+        """List available Splunk indexes.
+
+        Returns:
+            List of index names
+        """
+        response = self._request("GET", "/services/data/indexes", params={"output_mode": "json"})
+        return [entry["name"] for entry in response.get("entry", [])]
+
+    def search(
+        self,
+        query: str,
+        earliest_time: str = "-24h",
+        latest_time: str = "now",
+        max_count: int = 100,
+        output_mode: str = "json",
+    ) -> List[Dict[str, Any]]:
+        """Execute a Splunk search query (oneshot search for quick results).
+
+        Args:
+            query: SPL search query
+            earliest_time: Start time (e.g., "-24h", "2024-01-01T00:00:00")
+            latest_time: End time (e.g., "now", "2024-01-02T00:00:00")
+            max_count: Maximum number of results to return
+            output_mode: Output format (json, xml, csv)
+
+        Returns:
+            List of search results
+
+        Example:
+            >>> results = client.search(
+            ...     'index=main sourcetype=linux_secure "Failed password"',
+            ...     earliest_time="-1h",
+            ...     max_count=50
+            ... )
+        """
+        # Use oneshot search for quick results (no job creation)
+        data = {
+            "search": query if query.startswith("search") else f"search {query}",
+            "earliest_time": earliest_time,
+            "latest_time": latest_time,
+            "max_count": max_count,
+            "output_mode": output_mode,
+        }
+
+        response = self._request("POST", "/services/search/jobs/oneshot", data=data)
+
+        # Extract results from response
+        results = []
+        if "results" in response:
+            results = response["results"]
+        elif "entry" in response:
+            # Handle alternative response format
+            for entry in response["entry"]:
+                if "content" in entry:
+                    results.append(entry["content"])
+
+        return results
+
+    def create_search_job(self, query: str, earliest_time: str = "-24h", latest_time: str = "now", **kwargs: Any) -> str:
+        """Create an async search job for long-running queries.
+
+        Args:
+            query: SPL search query
+            earliest_time: Start time
+            latest_time: End time
+            **kwargs: Additional search parameters
+
+        Returns:
+            Search job ID (sid)
+
+        Example:
+            >>> sid = client.create_search_job(
+            ...     'index=* | stats count by sourcetype',
+            ...     earliest_time="-7d"
+            ... )
+            >>> results = client.get_search_results(sid)
+        """
+        data = {
+            "search": query if query.startswith("search") else f"search {query}",
+            "earliest_time": earliest_time,
+            "latest_time": latest_time,
+            "output_mode": "json",
+            **kwargs,
+        }
+
+        response = self._request("POST", "/services/search/jobs", data=data)
+
+        # Extract search ID from response
+        if "sid" in response:
+            return response["sid"]  # type: ignore[no-any-return]
+        elif "entry" in response and len(response["entry"]) > 0:
+            return response["entry"][0]["name"]  # type: ignore[no-any-return]
+
+        raise ValueError("Could not extract search job ID from response")
+
+    def get_search_job_status(self, sid: str) -> Dict[str, Any]:
+        """Get status of a search job.
+
+        Args:
+            sid: Search job ID
+
+        Returns:
+            Dict with job status information
+        """
+        return self._request("GET", f"/services/search/jobs/{sid}", params={"output_mode": "json"})  # type: ignore[no-any-return]
+
+    def wait_for_search_job(self, sid: str, poll_interval: int = 2, max_wait: int = 300) -> bool:
+        """Wait for search job to complete.
+
+        Args:
+            sid: Search job ID
+            poll_interval: Seconds between status checks
+            max_wait: Maximum seconds to wait
+
+        Returns:
+            True if job completed, False if timeout
+        """
+        elapsed = 0
+        while elapsed < max_wait:
+            status = self.get_search_job_status(sid)
+
+            # Check if job is done
+            if "entry" in status and len(status["entry"]) > 0:
+                content = status["entry"][0].get("content", {})
+                if content.get("isDone"):
+                    return True
+
+            time.sleep(poll_interval)
+            elapsed += poll_interval
+
+        return False
+
+    def get_search_results(
+        self, sid: str, offset: int = 0, count: int = 100, output_mode: str = "json"
+    ) -> List[Dict[str, Any]]:
+        """Get results from a completed search job.
+
+        Args:
+            sid: Search job ID
+            offset: Result offset (for pagination)
+            count: Number of results to return
+            output_mode: Output format
+
+        Returns:
+            List of search results
+        """
+        params = {
+            "output_mode": output_mode,
+            "offset": offset,
+            "count": count,
+        }
+
+        response = self._request("GET", f"/services/search/jobs/{sid}/results", params=params)
+
+        results = []
+        if "results" in response:
+            results = response["results"]
+        elif "entry" in response:
+            for entry in response["entry"]:
+                if "content" in entry:
+                    results.append(entry["content"])
+
+        return results
+
+    def delete_search_job(self, sid: str) -> None:
+        """Delete a search job.
+
+        Args:
+            sid: Search job ID
+        """
+        self._request("DELETE", f"/services/search/jobs/{sid}")
+
+    def search_async(
+        self,
+        query: str,
+        earliest_time: str = "-24h",
+        latest_time: str = "now",
+        max_results: int = 100,
+        wait: bool = True,
+        max_wait: int = 300,
+    ) -> List[Dict[str, Any]]:
+        """Execute a search asynchronously and return results.
+
+        This is useful for longer-running queries that may timeout with oneshot.
+
+        Args:
+            query: SPL search query
+            earliest_time: Start time
+            latest_time: End time
+            max_results: Maximum results to return
+            wait: Whether to wait for job completion
+            max_wait: Maximum seconds to wait for job
+
+        Returns:
+            List of search results
+
+        Example:
+            >>> results = client.search_async(
+            ...     'index=* | stats count by sourcetype',
+            ...     earliest_time="-7d",
+            ...     max_results=1000
+            ... )
+        """
+        # Create search job
+        sid = self.create_search_job(query, earliest_time, latest_time)
+
+        try:
+            if wait:
+                # Wait for completion
+                if not self.wait_for_search_job(sid, max_wait=max_wait):
+                    raise TimeoutError(f"Search job {sid} did not complete within {max_wait}s")
+
+            # Get results
+            return self.get_search_results(sid, count=max_results)
+
+        finally:
+            # Clean up search job
+            try:
+                self.delete_search_job(sid)
+            except Exception:
+                pass  # Ignore cleanup errors
+
+
+def create_client_from_env() -> SplunkClient:
+    """Create Splunk client from environment variables.
+
+    Environment variables:
+        SPLUNK_HOST: Splunk host
+        SPLUNK_TOKEN: Authentication token
+        SPLUNK_VERIFY_SSL: Whether to verify SSL (default: true)
+
+    Returns:
+        Configured SplunkClient instance
+
+    Raises:
+        ValueError: If required environment variables are missing
+    """
+    import os
+
+    host = os.getenv("SPLUNK_HOST")
+    token = os.getenv("SPLUNK_TOKEN")
+
+    if not host:
+        raise ValueError("SPLUNK_HOST environment variable is required")
+    if not token:
+        raise ValueError("SPLUNK_TOKEN environment variable is required")
+
+    verify_ssl = os.getenv("SPLUNK_VERIFY_SSL", "true").lower() in ("true", "1", "yes")
+
+    return SplunkClient(host=host, token=token, verify_ssl=verify_ssl)