#!/usr/bin/env python3
"""Agent for threat modeling with OWASP Threat Dragon.

Programmatically creates Threat Dragon JSON threat models, applies
STRIDE analysis to DFD elements, manages threat inventory, and
generates summary reports for security design reviews.
"""

import json
import sys
import uuid
from datetime import datetime


STRIDE_BY_ELEMENT = {
    "process": ["Spoofing", "Tampering", "Repudiation",
                "Information Disclosure", "Denial of Service",
                "Elevation of Privilege"],
    "data_store": ["Tampering", "Information Disclosure",
                   "Denial of Service"],
    "data_flow": ["Tampering", "Information Disclosure",
                  "Denial of Service"],
    "external_entity": ["Spoofing", "Repudiation"],
}

STRIDE_MITIGATIONS = {
    "Spoofing": ["Implement strong authentication (MFA)",
                 "Use mutual TLS for service-to-service"],
    "Tampering": ["Use integrity checks (HMAC, digital signatures)",
                  "Implement input validation"],
    "Repudiation": ["Enable comprehensive audit logging",
                    "Use tamper-evident log storage"],
    "Information Disclosure": ["Encrypt data at rest and in transit",
                               "Implement least-privilege access"],
    "Denial of Service": ["Implement rate limiting",
                          "Use auto-scaling and circuit breakers"],
    "Elevation of Privilege": ["Enforce RBAC and least privilege",
                               "Validate authorization on every request"],
}


class ThreatModelAgent:
    """Creates and manages OWASP Threat Dragon threat models."""

    def __init__(self, title, owner="Security Team", description=""):
        self.model = {
            "version": "2.2.0",
            "summary": {
                "title": title,
                "owner": owner,
                "description": description,
                "id": 0,
            },
            "detail": {
                "contributors": [],
                "diagrams": [],
                "diagramTop": 0,
                "reviewer": "",
                "threatTop": 0,
            },
        }
        self.threats = []
        self.threat_counter = 0

    def add_diagram(self, title, diagram_type="STRIDE"):
        """Add a new data flow diagram to the threat model."""
        diagram_id = len(self.model["detail"]["diagrams"])
        diagram = {
            "id": diagram_id,
            "title": title,
            "diagramType": diagram_type,
            "placeholder": f"New {diagram_type} diagram",
            "thumbnail": "",
            "version": "2.2.0",
            "cells": [],
        }
        self.model["detail"]["diagrams"].append(diagram)
        return diagram_id

    def add_element(self, diagram_id, element_type, name,
                    x=100, y=100, description=""):
        """Add a DFD element to a diagram."""
        element_id = str(uuid.uuid4())
        type_map = {
            "process": "tm.Process",
            "data_store": "tm.Store",
            "data_flow": "tm.Flow",
            "external_entity": "tm.Actor",
            "trust_boundary": "tm.Boundary",
        }
        cell = {
            "type": type_map.get(element_type, "tm.Process"),
            "id": element_id,
            "name": name,
            "description": description,
            "position": {"x": x, "y": y},
            "size": {"width": 100, "height": 60},
            "threats": [],
            "hasOpenThreats": False,
        }
        self.model["detail"]["diagrams"][diagram_id]["cells"].append(cell)
        return element_id

    def apply_stride(self, diagram_id, element_id, element_type):
        """Apply STRIDE analysis to a DFD element and generate threats."""
        categories = STRIDE_BY_ELEMENT.get(element_type, [])
        generated = []
        diagram = self.model["detail"]["diagrams"][diagram_id]
        element = next((c for c in diagram["cells"] if c["id"] == element_id), None)
        if not element:
            return []

        for category in categories:
            self.threat_counter += 1
            threat = {
                "id": str(self.threat_counter),
                "title": f"{category} - {element['name']}",
                "type": category,
                "status": "Open",
                "severity": "Medium",
                "description": f"Potential {category.lower()} threat "
                               f"against {element['name']}",
                "mitigation": "; ".join(
                    STRIDE_MITIGATIONS.get(category, ["Review required"])),
                "modelType": "STRIDE",
                "element_id": element_id,
            }
            element["threats"].append(threat)
            element["hasOpenThreats"] = True
            self.threats.append(threat)
            generated.append(threat)

        self.model["detail"]["threatTop"] = self.threat_counter
        return generated

    def update_threat_status(self, threat_id, status, mitigation=None):
        """Update a threat's status (Open, Mitigated, Not Applicable)."""
        for threat in self.threats:
            if threat["id"] == str(threat_id):
                threat["status"] = status
                if mitigation:
                    threat["mitigation"] = mitigation
                return threat
        return None

    def get_threat_summary(self):
        """Summarize threats by status and category."""
        summary = {"total": len(self.threats), "by_status": {},
                   "by_type": {}, "by_severity": {}}
        for t in self.threats:
            summary["by_status"][t["status"]] = \
                summary["by_status"].get(t["status"], 0) + 1
            summary["by_type"][t["type"]] = \
                summary["by_type"].get(t["type"], 0) + 1
            summary["by_severity"][t["severity"]] = \
                summary["by_severity"].get(t["severity"], 0) + 1
        return summary

    def save_model(self, output_path):
        """Save the threat model as Threat Dragon JSON file."""
        with open(output_path, "w") as f:
            json.dump(self.model, f, indent=2)
        return output_path

    def generate_report(self):
        """Generate threat model assessment report."""
        summary = self.get_threat_summary()
        report = {
            "title": self.model["summary"]["title"],
            "owner": self.model["summary"]["owner"],
            "report_date": datetime.utcnow().isoformat(),
            "diagrams": len(self.model["detail"]["diagrams"]),
            "threat_summary": summary,
            "open_threats": [t for t in self.threats if t["status"] == "Open"],
            "mitigated_threats": [t for t in self.threats
                                  if t["status"] == "Mitigated"],
        }
        print(json.dumps(report, indent=2))
        return report


def main():
    title = sys.argv[1] if len(sys.argv) > 1 else "Sample Application"
    output = sys.argv[2] if len(sys.argv) > 2 else "./threat_model.json"

    agent = ThreatModelAgent(title, owner="Security Team",
                             description="Automated threat model")
    did = agent.add_diagram("Main Data Flow")
    web = agent.add_element(did, "external_entity", "Web Browser", 50, 50)
    api = agent.add_element(did, "process", "API Gateway", 250, 50)
    db = agent.add_element(did, "data_store", "Database", 450, 50)
    flow1 = agent.add_element(did, "data_flow", "HTTPS Request", 150, 100)

    agent.apply_stride(did, web, "external_entity")
    agent.apply_stride(did, api, "process")
    agent.apply_stride(did, db, "data_store")
    agent.apply_stride(did, flow1, "data_flow")

    agent.save_model(output)
    agent.generate_report()


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
OWASP Threat Dragon Model Analyzer

Parses Threat Dragon JSON threat model files and generates
summary statistics, coverage reports, and mitigation gap analysis.
"""

import json
import sys
import os
from collections import defaultdict
from datetime import datetime


def load_threat_model(filepath: str) -> dict:
    with open(filepath) as f:
        return json.load(f)


def extract_threats(model: dict) -> list:
    threats = []
    detail = model.get("detail", {})
    for diagram in detail.get("diagrams", []):
        diagram_title = diagram.get("title", "Untitled")
        for cell in diagram.get("cells", []):
            cell_data = cell.get("data", {})
            cell_threats = cell_data.get("threats", [])
            for threat in cell_threats:
                threats.append({
                    "diagram": diagram_title,
                    "element": cell_data.get("name", cell.get("id", "unknown")),
                    "element_type": cell_data.get("type", "unknown"),
                    "title": threat.get("title", ""),
                    "description": threat.get("description", ""),
                    "severity": threat.get("severity", "Unknown"),
                    "status": threat.get("status", "Open"),
                    "type": threat.get("type", ""),
                    "mitigation": threat.get("mitigation", ""),
                    "model_type": threat.get("modelType", "STRIDE"),
                })
    return threats


def analyze_coverage(threats: list) -> dict:
    coverage = {
        "total_threats": len(threats),
        "by_status": defaultdict(int),
        "by_severity": defaultdict(int),
        "by_type": defaultdict(int),
        "by_element_type": defaultdict(int),
        "mitigated_count": 0,
        "open_count": 0,
        "not_applicable_count": 0,
        "with_mitigation_text": 0,
    }

    for threat in threats:
        status = threat["status"]
        coverage["by_status"][status] += 1
        coverage["by_severity"][threat["severity"]] += 1
        coverage["by_type"][threat["type"]] += 1
        coverage["by_element_type"][threat["element_type"]] += 1

        if status.lower() == "mitigated":
            coverage["mitigated_count"] += 1
        elif status.lower() == "open":
            coverage["open_count"] += 1
        elif status.lower() in ("not applicable", "n/a"):
            coverage["not_applicable_count"] += 1

        if threat["mitigation"].strip():
            coverage["with_mitigation_text"] += 1

    coverage["by_status"] = dict(coverage["by_status"])
    coverage["by_severity"] = dict(coverage["by_severity"])
    coverage["by_type"] = dict(coverage["by_type"])
    coverage["by_element_type"] = dict(coverage["by_element_type"])
    return coverage


def identify_gaps(threats: list) -> list:
    gaps = []
    for threat in threats:
        if threat["status"].lower() == "open" and not threat["mitigation"].strip():
            gaps.append({
                "diagram": threat["diagram"],
                "element": threat["element"],
                "threat_title": threat["title"],
                "severity": threat["severity"],
                "type": threat["type"],
            })
    return sorted(gaps, key=lambda g: {"Critical": 0, "High": 1, "Medium": 2, "Low": 3}.get(g["severity"], 4))


def stride_coverage_check(threats: list) -> dict:
    stride_categories = {
        "Spoofing": False,
        "Tampering": False,
        "Repudiation": False,
        "Information disclosure": False,
        "Denial of service": False,
        "Elevation of privilege": False,
    }
    for threat in threats:
        threat_type = threat.get("type", "")
        for category in stride_categories:
            if category.lower() in threat_type.lower():
                stride_categories[category] = True
    return stride_categories


def print_report(model: dict, coverage: dict, gaps: list, stride: dict) -> None:
    summary = model.get("summary", {})
    print(f"\n{'='*60}")
    print(f"Threat Model Analysis Report")
    print(f"{'='*60}")
    print(f"Title: {summary.get('title', 'Unknown')}")
    print(f"Owner: {summary.get('owner', 'Unknown')}")
    print(f"Description: {summary.get('description', '')}")
    print(f"Generated: {datetime.utcnow().isoformat()}Z")

    diagrams = model.get("detail", {}).get("diagrams", [])
    print(f"\nDiagrams: {len(diagrams)}")
    for d in diagrams:
        print(f"  - {d.get('title', 'Untitled')} ({d.get('diagramType', 'Unknown')} type)")

    print(f"\nThreat Summary:")
    print(f"  Total threats: {coverage['total_threats']}")
    print(f"  Mitigated: {coverage['mitigated_count']}")
    print(f"  Open: {coverage['open_count']}")
    print(f"  Not Applicable: {coverage['not_applicable_count']}")
    print(f"  With mitigation documented: {coverage['with_mitigation_text']}")

    if coverage["total_threats"] > 0:
        mitigation_rate = coverage["mitigated_count"] / coverage["total_threats"] * 100
        print(f"  Mitigation rate: {mitigation_rate:.1f}%")

    print(f"\nBy Severity:")
    for sev in ["Critical", "High", "Medium", "Low", "Unknown"]:
        count = coverage["by_severity"].get(sev, 0)
        if count:
            print(f"  {sev:12s}: {count}")

    print(f"\nSTRIDE Coverage:")
    for category, covered in stride.items():
        status = "COVERED" if covered else "MISSING"
        print(f"  {category:25s}: {status}")

    if gaps:
        print(f"\nMitigation Gaps ({len(gaps)} open threats without mitigations):")
        for gap in gaps:
            print(f"  [{gap['severity']}] {gap['threat_title']}")
            print(f"    Element: {gap['element']} | Diagram: {gap['diagram']}")
    else:
        print(f"\nNo mitigation gaps found.")


def main():
    if len(sys.argv) < 2:
        print("Usage: python process.py <threat_model.json>")
        print("  Analyzes an OWASP Threat Dragon JSON threat model file")
        sys.exit(1)

    filepath = sys.argv[1]
    if not os.path.exists(filepath):
        print(f"File not found: {filepath}")
        sys.exit(1)

    model = load_threat_model(filepath)
    threats = extract_threats(model)
    coverage = analyze_coverage(threats)
    gaps = identify_gaps(threats)
    stride = stride_coverage_check(threats)
    print_report(model, coverage, gaps, stride)

    output = filepath.replace(".json", "_analysis.json")
    analysis = {
        "model_title": model.get("summary", {}).get("title"),
        "analysis_date": datetime.utcnow().isoformat() + "Z",
        "coverage": coverage,
        "gaps": gaps,
        "stride_coverage": stride,
    }
    with open(output, "w") as f:
        json.dump(analysis, f, indent=2)
    print(f"\nAnalysis saved to: {output}")


if __name__ == "__main__":
    main()