07 - AI-Assisted Detection Generation: LLMs for Sigma Rule Authoring

# Base architecture of the AI-assisted pipeline
# File: ai_sigma_pipeline.py

import openai
import yaml
from pathlib import Path
from dataclasses import dataclass
from typing import Optional

@dataclass
class ThreatReport:
    content: str
    source_url: str
    report_type: str  # 'cti_blog', 'advisory', 'malware_analysis'

@dataclass
class GeneratedRule:
    sigma_yaml: str
    mitre_techniques: list[str]
    confidence: float
    validation_passed: bool
    test_results: Optional[dict] = None

class AISigmaPipeline:
    def __init__(self, openai_api_key: str, rules_output_dir: str):
        self.client = openai.OpenAI(api_key=openai_api_key)
        self.output_dir = Path(rules_output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def process_report(self, report: ThreatReport) -> list[GeneratedRule]:
        """Full pipeline from report to validated rules."""
        techniques = self._extract_attack_techniques(report)
        behaviors = self._extract_behaviors(report)

        rules = []
        for technique in techniques:
            rule = self._generate_sigma_rule(
                report=report,
                technique=technique,
                behaviors=behaviors
            )
            if rule:
                rules.append(rule)

        return [self._validate_and_test(r) for r in rules]

    def _extract_attack_techniques(self, report: ThreatReport) -> list[str]:
        """Extracts ATT&CK techniques from the report via LLM."""
        response = self.client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {
                    "role": "system",
                    "content": (
                        "You are a threat intelligence analyst expert in MITRE ATT&CK. "
                        "Extract ONLY ATT&CK techniques (format T1234 or T1234.001) "
                        "explicitly described in the text. Reply only with a JSON list."
                    )
                },
                {
                    "role": "user",
                    "content": f"Report:\n{report.content[:4000]}"
                }
            ],
            temperature=0.1
        )

        import json
        try:
            return json.loads(response.choices[0].message.content)
        except json.JSONDecodeError:
            return []

# Optimized system prompt for Sigma rule generation
SIGMA_SYSTEM_PROMPT = """
You are a Detection Engineer expert in writing Sigma rules.
When generating a Sigma rule, ALWAYS follow this YAML structure:

title: [descriptive title, max 80 chars]
id: [UUID v4]
status: experimental
description: [detailed description of the detected behavior]
author: AI-Assisted Detection
date: [today's date YYYY-MM-DD]
tags:
  - attack.[tactic]
  - attack.[technique]
logsource:
  category: [process_creation | network_connection | file_event | registry_event]
  product: [windows | linux | macos]
detection:
  [selector_name]:
    [field]: [value or list of values]
  condition: [selector_name]
falsepositives:
  - [plausible legitimate cases]
level: [informational | low | medium | high | critical]

CRITICAL RULES:
- ALWAYS use wildcards (*) in string values to avoid brittle exact matches
- Prefer fields with high availability (Image, CommandLine, ParentImage)
- Always indicate at least one realistic false positive
"""

def generate_with_cot(self, technique_id: str, report: ThreatReport) -> str:
    """Generation with explicit Chain-of-Thought."""

    # Step 1: Ask the model to reason first
    reasoning_response = self.client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": SIGMA_SYSTEM_PROMPT},
            {
                "role": "user",
                "content": f"""Before writing the rule for {technique_id}, analyze:
1. What forensic artifacts does this technique leave in logs?
2. What is the most appropriate logsource?
3. Which fields have the highest signal/noise discrimination?
4. What are the most common false positives?

Report: {report.content[:2000]}"""
            }
        ],
        temperature=0.3
    )

    reasoning = reasoning_response.choices[0].message.content

    # Step 2: Use the reasoning to guide generation
    rule_response = self.client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": SIGMA_SYSTEM_PROMPT},
            {"role": "user", "content": f"Technical analysis:\n{reasoning}"},
            {
                "role": "assistant",
                "content": "Based on this analysis, I will generate the optimal Sigma rule:"
            },
            {"role": "user", "content": "Proceed with YAML generation."}
        ],
        temperature=0.1
    )

    return rule_response.choices[0].message.content

import yaml
from sigma.rule import SigmaRule
from sigma.exceptions import SigmaError
import re, uuid

class SigmaRuleValidator:
    VALID_CATEGORIES = {
        'process_creation', 'network_connection', 'file_event',
        'registry_event', 'dns_query', 'image_load', 'pipe_created'
    }
    VALID_LEVELS = {'informational', 'low', 'medium', 'high', 'critical'}

    def validate(self, sigma_yaml: str) -> tuple[bool, list[str]]:
        """Validates a Sigma rule. Returns (valid, errors)."""
        errors = []

        try:
            rule_dict = yaml.safe_load(sigma_yaml)
        except yaml.YAMLError as e:
            return False, [f"Invalid YAML: {str(e)}"]

        # Required fields check
        for field in ['title', 'description', 'logsource', 'detection']:
            if field not in rule_dict:
                errors.append(f"Missing required field: {field}")

        if errors:
            return False, errors

        # Logsource validation
        logsource = rule_dict.get('logsource', {})
        if 'category' in logsource:
            if logsource['category'] not in self.VALID_CATEGORIES:
                errors.append(f"Invalid logsource category: {logsource['category']}")

        # Detection/condition validation
        detection = rule_dict.get('detection', {})
        if 'condition' not in detection:
            errors.append("Missing 'condition' field in detection")

        # UUID validation
        rule_id = rule_dict.get('id', '')
        if rule_id:
            try:
                uuid.UUID(str(rule_id))
            except ValueError:
                errors.append(f"ID is not a valid UUID: {rule_id}")

        # pySigma deep validation
        try:
            SigmaRule.from_yaml(sigma_yaml)
        except SigmaError as e:
            errors.append(f"pySigma error: {str(e)}")

        return len(errors) == 0, errors

class SigmaRuleTester:
    def generate_test_events(self, sigma_yaml: str, llm_client) -> dict:
        """Generates test events via LLM based on the rule."""
        prompt = f"""Given this Sigma rule:
{sigma_yaml}

Generate in JSON format two event lists:
1. "true_positives": 3 events that MUST trigger the rule
2. "false_positives": 3 legitimate events that must NOT trigger

Format:
{
  "true_positives": [{"Image": "...", "CommandLine": "..."}],
  "false_positives": [{"Image": "...", "CommandLine": "..."}]
}"""

        response = llm_client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.2
        )

        import json
        try:
            return json.loads(response.choices[0].message.content)
        except json.JSONDecodeError:
            return {"true_positives": [], "false_positives": []}

# GitHub Actions for AI-generated rules
# .github/workflows/ai-sigma-generation.yml
"""
name: AI Sigma Rule Generation
on:
  schedule:
    - cron: '0 6 * * *'
jobs:
  generate-rules:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - name: Setup Python
        uses: actions/setup-python@v5
        with: { python-version: '3.12' }
      - name: Run AI pipeline
        env:
          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
        run: python scripts/ai_sigma_pipeline.py
      - name: Create Pull Request
        uses: peter-evans/create-pull-request@v6
        with:
          title: '[AI-Generated] Detection rules from CTI feed'
          labels: 'ai-generated,needs-review'
"""