Arkime Integration

#!/usr/bin/env python3
"""
Arkime REST API client for programmatic session search and PCAP retrieval.
"""
import requests
import json
from datetime import datetime, timezone
from urllib.parse import urlencode
from pathlib import Path

ARKIME_URL = "http://arkime:8005"
ARKIME_USER = "admin"
ARKIME_PASS = "your_password"

session = requests.Session()
session.auth = (ARKIME_USER, ARKIME_PASS)
session.verify = False  # Set to True with proper TLS cert


def search_sessions(query: str, start_time: datetime, stop_time: datetime,
                    limit: int = 100) -> dict:
    """Search Arkime sessions by query expression and time range."""
    params = {
        "expression": query,
        "startTime": int(start_time.timestamp()),
        "stopTime": int(stop_time.timestamp()),
        "length": limit,
        "fields": "ip.src,ip.dst,port.src,port.dst,starttime,stoptime,"
                  "bytes.src,bytes.dst,tls.ja3,tls.sni,http.uri",
    }
    url = f"{ARKIME_URL}/api/sessions"
    resp = session.get(url, params=params, timeout=30)
    resp.raise_for_status()
    return resp.json()


def get_session_pcap(session_id: str, node: str, output_path: str) -> int:
    """Download PCAP for a specific session. Returns file size in bytes."""
    url = f"{ARKIME_URL}/{node}/session/{session_id}/pcap"
    resp = session.get(url, timeout=60, stream=True)
    resp.raise_for_status()

    path = Path(output_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    total = 0
    with open(path, 'wb') as f:
        for chunk in resp.iter_content(chunk_size=65536):
            f.write(chunk)
            total += len(chunk)
    return total


def tag_session(session_id: str, node: str, tags: list) -> bool:
    """Add tags to a session (for analyst notes, threat intel labels)."""
    url = f"{ARKIME_URL}/api/sessions/addtags"
    data = {
        "ids": f"{node},{session_id}",
        "tags": ",".join(tags),
    }
    resp = session.post(url, json=data, timeout=10)
    return resp.ok


def find_and_retrieve_pcap_for_alert(
    src_ip: str, dst_ip: str, dst_port: int,
    alert_time: datetime, output_dir: str = "/tmp/arkime-pcaps"
) -> list:
    """
    Given an alert's 5-tuple, find matching Arkime sessions and download PCAPs.
    Returns list of downloaded file paths.
    """
    from datetime import timedelta

    # Search window: ±5 minutes around the alert time
    start = alert_time - timedelta(minutes=5)
    stop = alert_time + timedelta(minutes=5)

    query = (f"ip.src=={src_ip} && ip.dst=={dst_ip} && port.dst=={dst_port}"
             f" && protocols==tls")

    print(f"Searching Arkime: {query}")
    results = search_sessions(query, start, stop)

    sessions_data = results.get("data", [])
    print(f"Found {len(sessions_data)} matching sessions")

    downloaded = []
    for s in sessions_data:
        sid = s.get("id")
        node = s.get("node", "default")
        src = s.get("source", {})
        ts = s.get("starttime", "")

        output_file = f"{output_dir}/{src_ip}_{dst_ip}_{dst_port}_{sid[:8]}.pcap"
        print(f"Downloading session {sid[:8]} ({ts})...")

        size = get_session_pcap(sid, node, output_file)
        print(f"  Saved {size} bytes to {output_file}")

        # Tag session as investigated
        tag_session(sid, node, ["investigated", "auto-retrieved"])

        downloaded.append(output_file)

    return downloaded


# Example: triggered by a Suricata alert
if __name__ == "__main__":
    alert_time = datetime(2024, 1, 15, 14, 23, 17, tzinfo=timezone.utc)
    pcaps = find_and_retrieve_pcap_for_alert(
        src_ip="10.0.0.50",
        dst_ip="1.2.3.4",
        dst_port=443,
        alert_time=alert_time,
        output_dir="/tmp/incident-20240115"
    )
    print(f"\nDownloaded {len(pcaps)} PCAPs:")
    for p in pcaps:
        print(f"  {p}")