# SecuBox Threat Analyst - Rule Generators
# Generates filters for mitmproxy, CrowdSec, WAF

# =============================================================================
# MITMPROXY FILTER GENERATION
# =============================================================================

generate_mitmproxy_filters() {
	local analysis="$1"
	local threats="$2"

	# Extract IPs to block from threats
	local blocked_ips=$(echo "$threats" | jsonfilter -e '@[*].ip' 2>/dev/null | sort -u | head -20)

	# Extract URL patterns from mitmproxy threats
	local url_patterns=$(echo "$threats" | jsonfilter -e '@[?(@.source=="mitmproxy")].url' 2>/dev/null | \
		grep -oE '/[a-zA-Z0-9_/-]+' | sort | uniq -c | sort -rn | head -10 | awk '{print $2}')

	# Build Python filter
	cat <<'PYTHON_HEADER'
# SecuBox AI-Generated mitmproxy Filters
# Auto-generated by Threat Analyst Agent
# Do not edit manually - will be overwritten

import re
from mitmproxy import http, ctx

class AIThreatFilter:
    """AI-generated threat detection filters"""

    def __init__(self):
        self.blocked_ips = set([
PYTHON_HEADER

	# Add blocked IPs
	for ip in $blocked_ips; do
		[ -n "$ip" ] && echo "            \"$ip\","
	done

	cat <<'PYTHON_MIDDLE'
        ])

        self.suspicious_patterns = [
PYTHON_MIDDLE

	# Add URL patterns
	for pattern in $url_patterns; do
		[ -n "$pattern" ] && echo "            r\"$(echo "$pattern" | sed 's/\\/\\\\/g')\","
	done

	# Add common attack patterns
	cat <<'PYTHON_PATTERNS'
            r"/wp-admin",
            r"/\.env",
            r"/\.git",
            r"/phpinfo",
            r"/eval\(",
            r"/base64_decode",
            r"/<script",
            r"/union\s+select",
            r"/etc/passwd",
            r"/proc/self",
        ]

        self.blocked_user_agents = [
            r"sqlmap",
            r"nikto",
            r"nmap",
            r"masscan",
            r"zgrab",
            r"censys",
            r"shodan",
        ]

    def request(self, flow: http.HTTPFlow) -> None:
        client_ip = flow.client_conn.peername[0]

        # Block known bad IPs
        if client_ip in self.blocked_ips:
            flow.kill()
            ctx.log.warn(f"[AI-FILTER] Blocked IP: {client_ip}")
            return

        # Check URL patterns
        url = flow.request.pretty_url
        for pattern in self.suspicious_patterns:
            if re.search(pattern, url, re.IGNORECASE):
                self._log_threat(flow, "suspicious_url", pattern)
                flow.kill()
                return

        # Check User-Agent
        ua = flow.request.headers.get("User-Agent", "")
        for pattern in self.blocked_user_agents:
            if re.search(pattern, ua, re.IGNORECASE):
                self._log_threat(flow, "bad_useragent", pattern)
                flow.kill()
                return

    def _log_threat(self, flow, threat_type, pattern):
        import json
        import time
        log_entry = {
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
            "client_ip": flow.client_conn.peername[0],
            "url": flow.request.pretty_url,
            "threat_type": threat_type,
            "pattern": pattern,
            "source": "ai_filter"
        }
        with open("/srv/mitmproxy/ai_threats.log", "a") as f:
            f.write(json.dumps(log_entry) + "\n")
        ctx.log.warn(f"[AI-FILTER] {threat_type}: {pattern}")

addons = [AIThreatFilter()]
PYTHON_PATTERNS
}

# =============================================================================
# CROWDSEC SCENARIO GENERATION
# =============================================================================

generate_crowdsec_scenario() {
	local analysis="$1"
	local threats="$2"

	# Count threat types
	local threat_types=$(echo "$threats" | jsonfilter -e '@[*].type' 2>/dev/null | sort | uniq -c | sort -rn)

	# Extract top scenarios from CrowdSec alerts
	local top_scenarios=$(echo "$threats" | jsonfilter -e '@[?(@.source=="crowdsec")].scenario' 2>/dev/null | \
		sort | uniq -c | sort -rn | head -5)

	cat <<YAML_HEADER
# SecuBox AI-Generated CrowdSec Scenario
# Auto-generated by Threat Analyst Agent
# Location: /etc/crowdsec/scenarios/ai-generated.yaml

type: leaky
name: secubox/ai-threat-pattern
description: "AI-detected threat patterns from SecuBox analysis"
filter: "evt.Meta.log_type in ['mitmproxy_threat', 'netifyd_risk', 'auth_failure']"
groupby: evt.Meta.source_ip
capacity: 3
leakspeed: 30s
blackhole: 5m
labels:
  service: secubox
  type: ai_detection
  remediation: true

# Trigger conditions based on analyzed patterns
YAML_HEADER

	# Add dynamic conditions based on threat analysis
	cat <<YAML_CONDITIONS
# Attack pattern conditions (AI-generated)
---
type: conditional
name: secubox/ai-scan-detection
description: "AI-detected scanning behavior"
filter: |
  evt.Meta.log_type == 'mitmproxy_threat'
  and evt.Meta.threat_type in ['scan', 'probe', 'enumeration']
groupby: evt.Meta.source_ip
condition: |
  len(evt.GetSources()) > 5
labels:
  service: secubox
  confidence: high

---
type: conditional
name: secubox/ai-injection-attempt
description: "AI-detected injection attempts"
filter: |
  evt.Meta.log_type == 'mitmproxy_threat'
  and evt.Meta.threat_type contains 'injection'
groupby: evt.Meta.source_ip
condition: |
  len(evt.GetSources()) >= 2
labels:
  service: secubox
  confidence: high
  cve: potential
YAML_CONDITIONS
}

# =============================================================================
# WAF RULES GENERATION
# =============================================================================

generate_waf_rules() {
	local analysis="$1"
	local threats="$2"

	# Extract patterns from threats
	local attack_urls=$(echo "$threats" | jsonfilter -e '@[?(@.source=="mitmproxy")].url' 2>/dev/null | head -50)

	cat <<JSON_HEADER
{
  "_comment": "SecuBox AI-Generated WAF Rules",
  "_generated": "$(date -Iseconds)",
  "_source": "threat-analyst",
  "rules": [
JSON_HEADER

	local first=1

	# SQL Injection patterns
	[ $first -eq 0 ] && echo ","
	first=0
	cat <<'RULE_SQLI'
    {
      "id": "ai-sqli-001",
      "name": "AI-Detected SQL Injection",
      "severity": "critical",
      "patterns": [
        "union\\s+select",
        "select\\s+.*\\s+from",
        "insert\\s+into",
        "drop\\s+table",
        "or\\s+1\\s*=\\s*1",
        "and\\s+1\\s*=\\s*1",
        "sleep\\s*\\(",
        "benchmark\\s*\\("
      ],
      "action": "block",
      "log": true
    }
RULE_SQLI

	# XSS patterns
	echo ","
	cat <<'RULE_XSS'
    {
      "id": "ai-xss-001",
      "name": "AI-Detected XSS Attempt",
      "severity": "high",
      "patterns": [
        "<script[^>]*>",
        "javascript:",
        "onerror\\s*=",
        "onload\\s*=",
        "onclick\\s*=",
        "eval\\s*\\(",
        "document\\.cookie"
      ],
      "action": "block",
      "log": true
    }
RULE_XSS

	# Path traversal
	echo ","
	cat <<'RULE_TRAVERSAL'
    {
      "id": "ai-traversal-001",
      "name": "AI-Detected Path Traversal",
      "severity": "high",
      "patterns": [
        "\\.\\./",
        "\\.\\.\\\\/",
        "/etc/passwd",
        "/etc/shadow",
        "/proc/self",
        "file://"
      ],
      "action": "block",
      "log": true
    }
RULE_TRAVERSAL

	# Scanner detection
	echo ","
	cat <<'RULE_SCANNER'
    {
      "id": "ai-scanner-001",
      "name": "AI-Detected Scanner",
      "severity": "medium",
      "user_agent_patterns": [
        "sqlmap",
        "nikto",
        "nmap",
        "masscan",
        "zgrab",
        "censys",
        "shodan",
        "nuclei"
      ],
      "action": "block",
      "log": true
    }
RULE_SCANNER

	cat <<JSON_FOOTER
  ]
}
JSON_FOOTER
}

# =============================================================================
# QUEUE MANAGEMENT
# =============================================================================

queue_rule() {
	local rule_type="$1"
	local rule_content="$2"

	local pending_file="$STATE_DIR/pending_rules.json"
	local rule_id=$(date +%s)_${rule_type}

	# Create queue file if needed
	[ ! -f "$pending_file" ] && echo '[]' > "$pending_file"

	# Add rule to queue
	local new_rule=$(cat <<EOF
{"id":"$rule_id","type":"$rule_type","created":"$(date -Iseconds)","content":"$(echo "$rule_content" | base64 -w 0)"}
EOF
)

	# Append to queue (simple approach)
	local current=$(cat "$pending_file")
	echo "$current" | sed 's/]$//' > "$pending_file.tmp"
	[ "$current" != "[]" ] && echo "," >> "$pending_file.tmp"
	echo "$new_rule" >> "$pending_file.tmp"
	echo "]" >> "$pending_file.tmp"
	mv "$pending_file.tmp" "$pending_file"

	log_info "Queued rule $rule_id ($rule_type)"
}
