marcredhat-siem-toolkit-patched/backend/services/s1_client.py

import os
import asyncio
import httpx
from datetime import datetime, timezone

BASE_URL = os.environ.get("S1_BASE_URL", "https://demo.sentinelone.net").rstrip("/")
TOKEN = os.environ.get("S1_API_TOKEN", "")

# Scalyr/XDR PowerQuery credentials — from SDL_XDR_URL + SDL_LOG_READ_KEY
# in the SentinelOne console: Settings → Integrations → Data Lake API Keys
SDL_XDR_URL = os.environ.get("SDL_XDR_URL", "https://xdr.us1.sentinelone.net").rstrip("/")
SDL_LOG_READ_KEY = os.environ.get("SDL_LOG_READ_KEY", "")

# Management Console API uses ApiToken auth
HEADERS = {
    "Authorization": f"ApiToken {TOKEN}",
    "Content-Type": "application/json",
}


def _iso_to_epoch_ms(iso_str: str) -> int:
    """Convert ISO-8601 UTC string to epoch milliseconds for Scalyr API."""
    dt = datetime.fromisoformat(iso_str.replace("Z", "+00:00"))
    return int(dt.timestamp() * 1000)


async def get_star_rules(limit: int = 200) -> list:
    """Fetch active STAR rules from the Management Console API."""
    async with httpx.AsyncClient(timeout=30) as client:
        resp = await client.get(
            f"{BASE_URL}/web/api/v2.1/cloud-detection/rules",
            headers=HEADERS,
            params={"limit": limit},
        )
        resp.raise_for_status()
        return resp.json().get("data", [])


async def run_powerquery(query: str, from_date: str, to_date: str) -> dict:
    """
    Run a PowerQuery against the Singularity Data Lake via the Scalyr XDR API.
    Uses SDL_XDR_URL + SDL_LOG_READ_KEY (Scalyr readlog token).
    The Scalyr PowerQuery API is synchronous — results return in one request.
    """
    if not SDL_LOG_READ_KEY:
        return {"events": [], "error": "SDL_LOG_READ_KEY not configured — add it to .env"}

    start_ms = _iso_to_epoch_ms(from_date)
    end_ms = _iso_to_epoch_ms(to_date)

    payload = {
        "token": SDL_LOG_READ_KEY,
        "query": query,
        "startTime": start_ms,
        "endTime": end_ms,
        "maxCount": 1000,
    }

    async with httpx.AsyncClient(timeout=120) as client:
        for attempt in range(3):
            try:
                resp = await client.post(
                    f"{SDL_XDR_URL}/api/powerQuery",
                    json=payload,
                )
                resp.raise_for_status()
                break
            except httpx.HTTPStatusError as e:
                if e.response.status_code == 429 and attempt < 2:
                    await asyncio.sleep(10 * (attempt + 1))
                    continue
                raise RuntimeError(
                    f"HTTP {e.response.status_code} from {e.request.url}: {e.response.text[:500]}"
                ) from e

        data = resp.json()
        status = data.get("status", "")

        if status != "success":
            # Return full response as error detail for debugging
            return {"events": [], "error": f"PowerQuery status={status}: {str(data)[:400]}"}

        # Scalyr PowerQuery returns: {"status":"success","columns":[{"name":"..."},...], "values":[[...],...],...}
        raw_cols = data.get("columns", [])
        values = data.get("values", [])

        if raw_cols and values:
            # columns may be list of strings or list of {"name":...} dicts
            col_names = [
                c["name"] if isinstance(c, dict) else c
                for c in raw_cols
            ]
            rows = [dict(zip(col_names, row)) for row in values]
            return {"events": rows}

        # Fallback: return raw matches array
        matches = data.get("matches", [])
        return {"events": matches}


async def list_sdl_parsers() -> list[str]:
    """List all parser filenames under /logParsers/ in SDL."""
    async with httpx.AsyncClient(timeout=30) as client:
        resp = await client.get(
            f"{BASE_URL}/api/v1/files/logParsers",
            headers=HEADERS,
        )
        resp.raise_for_status()
        data = resp.json()
        # Response is a list of file objects or a dict with 'files' key
        if isinstance(data, list):
            return [f.get("name") or f.get("path", "") for f in data if isinstance(f, dict)]
        return [f.get("name") or f.get("path", "") for f in data.get("files", [])]


async def get_sdl_parser(filename: str) -> dict:
    """Fetch a single SDL parser file by name."""
    async with httpx.AsyncClient(timeout=30) as client:
        resp = await client.get(
            f"{BASE_URL}/api/v1/files/logParsers/{filename}",
            headers=HEADERS,
        )
        resp.raise_for_status()
        return resp.json()


async def get_sites() -> list:
    async with httpx.AsyncClient(timeout=30) as client:
        resp = await client.get(
            f"{BASE_URL}/web/api/v2.1/sites",
            headers=HEADERS,
            params={"limit": 100},
        )
        resp.raise_for_status()
        return resp.json().get("data", {}).get("sites", [])