icdev 0.0.3__py3-none-any.whl

tools/audit_engine/self_heal.py

@@ -0,0 +1,1042 @@
+#!/usr/bin/env python3
+# CUI // SP-CTI
+# Controlled by: Department of Defense
+# CUI Category: CTI
+# Distribution: D
+# POC: ICDEV System Administrator
+"""Audit Engine — Self-Heal Remediation Engine.
+
+Executes authorized remediation actions for audit findings.
+Each action is deterministic, reversible, and logged to audit trail.
+
+Actions require explicit user authorization before execution.
+"""
+
+import json
+import os
+import re
+import shutil
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Dict, List, Optional
+
+BASE_DIR = Path(__file__).resolve().parent.parent.parent
+
+# Safe import of storage layer
+try:
+    from tools.db.storage import get_connection
+except ImportError:
+    get_connection = None
+
+
+# ---------------------------------------------------------------
+# Remediation Action Registry
+# ---------------------------------------------------------------
+
+REMEDIATION_ACTIONS = {
+    "AC-6": {
+        "id": "AC-6",
+        "title": "Enforce least privilege",
+        "action_type": "config_change",
+        "reversible": True,
+        "description": "Add non-root container configuration and read-only filesystem settings",
+        "handler": "_remediate_least_privilege",
+    },
+    "AC-16": {
+        "id": "AC-16",
+        "title": "Add CUI classification markings",
+        "action_type": "code_annotation",
+        "reversible": True,
+        "description": "Add CUI // SP-CTI header markings to unmarked Python files",
+        "handler": "_remediate_cui_markings",
+    },
+    "SI-10": {
+        "id": "SI-10",
+        "title": "Add input validation patterns",
+        "action_type": "code_pattern",
+        "reversible": True,
+        "description": "Add input validation wrappers to identified unsafe patterns",
+        "handler": "_remediate_input_validation",
+    },
+    "AU-2": {
+        "id": "AU-2",
+        "title": "Enable audit logging",
+        "action_type": "code_pattern",
+        "reversible": True,
+        "description": "Add logging imports and basic audit logging to modules",
+        "handler": "_remediate_audit_logging",
+    },
+    "CM-6": {
+        "id": "CM-6",
+        "title": "Add security configuration files",
+        "action_type": "config_change",
+        "reversible": True,
+        "description": "Create .gitignore, SECURITY.md, and security configuration",
+        "handler": "_remediate_security_config",
+    },
+    "IA-5": {
+        "id": "IA-5",
+        "title": "Rotate detected secrets",
+        "action_type": "config_change",
+        "reversible": False,
+        "description": "Move hardcoded secrets to environment variables",
+        "handler": "_remediate_secrets",
+    },
+    "SA-11": {
+        "id": "SA-11",
+        "title": "Generate SBOM",
+        "action_type": "tool_enable",
+        "reversible": True,
+        "description": "Generate Software Bill of Materials for the project",
+        "handler": "_remediate_sbom",
+    },
+    "RA-5": {
+        "id": "RA-5",
+        "title": "Enable vulnerability scanning",
+        "action_type": "tool_enable",
+        "reversible": True,
+        "description": "Enable SAST and dependency scanning in CI/CD pipeline",
+        "handler": "_remediate_vuln_scanning",
+    },
+    "SI-2": {
+        "id": "SI-2",
+        "title": "Remediate known vulnerabilities",
+        "action_type": "dependency_update",
+        "reversible": True,
+        "description": "Update dependencies with known CVEs to patched versions",
+        "handler": "_remediate_dependencies",
+    },
+    "SC-13": {
+        "id": "SC-13",
+        "title": "Enable encryption",
+        "action_type": "config_change",
+        "reversible": True,
+        "description": "Add TLS/encryption configuration to communication channels",
+        "handler": "_remediate_encryption",
+    },
+}
+
+
+def list_available_actions(audit_result: Dict) -> List[Dict]:
+    """List remediation actions available for a given audit result.
+
+    Returns actions with their authorization status and estimated impact.
+    """
+    actions = []
+    if not audit_result or not audit_result.get("regime_scores"):
+        return actions
+
+    # Collect all failing controls across all regimes
+    failing_controls = set()
+    for regime_id, regime_data in audit_result["regime_scores"].items():
+        if regime_data.get("error"):
+            continue
+        for cat in regime_data.get("categories", []):
+            for ctrl in cat.get("controls", []):
+                if ctrl.get("status") in ("fail", "partial"):
+                    # Extract NIST family prefix
+                    ctrl_id = ctrl.get("control_id", "")
+                    family = _extract_control_family(ctrl_id)
+                    if family:
+                        failing_controls.add(family)
+
+    for family in failing_controls:
+        if family in REMEDIATION_ACTIONS:
+            action = REMEDIATION_ACTIONS[family].copy()
+            action["affected_regimes"] = _count_affected_regimes(
+                audit_result, family
+            )
+            actions.append(action)
+
+    # Sort by number of affected regimes (most impactful first)
+    actions.sort(key=lambda a: a.get("affected_regimes", 0), reverse=True)
+    return actions
+
+
+def execute_remediation(
+    control_id: str,
+    target_path: str,
+    project_id: Optional[str] = None,
+    dry_run: bool = False,
+) -> Dict:
+    """Execute a single remediation action.
+
+    Args:
+        control_id: NIST control family ID (e.g., 'AC-16', 'SI-10').
+        target_path: Directory to remediate.
+        project_id: Optional project ID.
+        dry_run: If True, preview changes without applying.
+
+    Returns:
+        Result dict with status, changes made, and before/after state.
+    """
+    target_path = os.path.abspath(target_path)
+    timestamp = datetime.now(timezone.utc).isoformat()
+
+    if control_id not in REMEDIATION_ACTIONS:
+        return {
+            "success": False,
+            "error": f"No remediation available for control {control_id}",
+            "control_id": control_id,
+        }
+
+    action = REMEDIATION_ACTIONS[control_id]
+    handler_name = action["handler"]
+    handler = globals().get(handler_name)
+
+    if not handler:
+        return {
+            "success": False,
+            "error": f"Handler {handler_name} not implemented",
+            "control_id": control_id,
+        }
+
+    result = {
+        "control_id": control_id,
+        "action_title": action["title"],
+        "action_type": action["action_type"],
+        "reversible": action["reversible"],
+        "target_path": target_path,
+        "dry_run": dry_run,
+        "timestamp": timestamp,
+        "changes": [],
+    }
+
+    try:
+        changes = handler(target_path, project_id, dry_run)
+        result["changes"] = changes
+        result["success"] = True
+        result["files_modified"] = len([c for c in changes if c.get("applied")])
+    except Exception as e:
+        result["success"] = False
+        result["error"] = str(e)
+
+    # Log to audit trail
+    if not dry_run and result["success"]:
+        _log_remediation(result)
+
+    return result
+
+
+# ---------------------------------------------------------------
+# Remediation Handlers
+# ---------------------------------------------------------------
+
+CUI_HEADER = """# CUI // SP-CTI
+# Controlled by: Department of Defense
+# CUI Category: CTI
+# Distribution: D
+# POC: ICDEV System Administrator
+"""
+
+
+def _remediate_cui_markings(
+    target_path: str, project_id: Optional[str], dry_run: bool
+) -> List[Dict]:
+    """Add CUI markings to unmarked Python files."""
+    changes = []
+    skip_dirs = {
+        ".git", "node_modules", "venv", "__pycache__", ".tmp", "data",
+        ".venv", "env", "dist", "build", "site-packages",
+    }
+    count = 0
+    max_files = 200
+
+    for root, dirs, files in os.walk(target_path):
+        dirs[:] = [d for d in dirs if d not in skip_dirs]
+        for fname in files:
+            if count >= max_files:
+                break
+            if not fname.endswith(".py"):
+                continue
+            fpath = os.path.join(root, fname)
+            try:
+                with open(fpath, "r", encoding="utf-8", errors="ignore") as f:
+                    content = f.read(500)
+                if "CUI" in content or "UNCLASSIFIED" in content:
+                    continue  # Already marked
+
+                rel_path = os.path.relpath(fpath, target_path)
+                change = {
+                    "file": rel_path,
+                    "action": "add_cui_header",
+                    "applied": not dry_run,
+                }
+
+                if not dry_run:
+                    with open(fpath, "r", encoding="utf-8") as f:
+                        full_content = f.read()
+                    # Preserve shebang if present
+                    if full_content.startswith("#!"):
+                        shebang_end = full_content.index("\n") + 1
+                        new_content = (
+                            full_content[:shebang_end] + CUI_HEADER + full_content[shebang_end:]
+                        )
+                    else:
+                        new_content = CUI_HEADER + full_content
+                    with open(fpath, "w", encoding="utf-8") as f:
+                        f.write(new_content)
+
+                changes.append(change)
+                count += 1
+            except (OSError, UnicodeDecodeError):
+                continue
+        if count >= max_files:
+            break
+
+    return changes
+
+
+def _remediate_audit_logging(
+    target_path: str, project_id: Optional[str], dry_run: bool
+) -> List[Dict]:
+    """Add logging imports to Python modules missing them."""
+    changes = []
+    skip_dirs = {
+        ".git", "node_modules", "venv", "__pycache__", ".tmp", "data",
+        ".venv", "env", "dist", "build", "site-packages",
+    }
+    count = 0
+    max_files = 50
+
+    for root, dirs, files in os.walk(target_path):
+        dirs[:] = [d for d in dirs if d not in skip_dirs]
+        for fname in files:
+            if count >= max_files:
+                break
+            if not fname.endswith(".py") or fname == "__init__.py":
+                continue
+            fpath = os.path.join(root, fname)
+            try:
+                with open(fpath, "r", encoding="utf-8", errors="ignore") as f:
+                    content = f.read()
+                if "import logging" in content or "from logging" in content:
+                    continue
+
+                # Only add to files with functions/classes
+                if "def " not in content and "class " not in content:
+                    continue
+
+                rel_path = os.path.relpath(fpath, target_path)
+                change = {
+                    "file": rel_path,
+                    "action": "add_logging_import",
+                    "applied": not dry_run,
+                }
+
+                if not dry_run:
+                    # Insert logging import after other imports
+                    lines = content.split("\n")
+                    insert_idx = 0
+                    for i, line in enumerate(lines):
+                        if line.startswith("import ") or line.startswith("from "):
+                            insert_idx = i + 1
+                    lines.insert(insert_idx, "import logging")
+                    lines.insert(
+                        insert_idx + 1,
+                        f'logger = logging.getLogger(__name__)',
+                    )
+                    lines.insert(insert_idx + 2, "")
+                    with open(fpath, "w", encoding="utf-8") as f:
+                        f.write("\n".join(lines))
+
+                changes.append(change)
+                count += 1
+            except (OSError, UnicodeDecodeError):
+                continue
+        if count >= max_files:
+            break
+
+    return changes
+
+
+def _remediate_security_config(
+    target_path: str, project_id: Optional[str], dry_run: bool
+) -> List[Dict]:
+    """Create missing security configuration files."""
+    changes = []
+
+    # .gitignore
+    gitignore_path = Path(target_path) / ".gitignore"
+    if not gitignore_path.exists():
+        change = {"file": ".gitignore", "action": "create", "applied": not dry_run}
+        if not dry_run:
+            gitignore_path.write_text(
+                "# Secrets\n.env\n.env.*\ncredentials.*\n*.pem\n*.key\nid_rsa*\n\n"
+                "# Build\ndist/\nbuild/\n*.egg-info/\n\n"
+                "# Python\n__pycache__/\n*.pyc\n.venv/\nvenv/\n\n"
+                "# Data\n*.db\ndata/\n\n"
+                "# IDE\n.vscode/\n.idea/\n",
+                encoding="utf-8",
+            )
+        changes.append(change)
+
+    # SECURITY.md
+    security_path = Path(target_path) / "SECURITY.md"
+    if not security_path.exists():
+        change = {"file": "SECURITY.md", "action": "create", "applied": not dry_run}
+        if not dry_run:
+            security_path.write_text(
+                "# Security Policy\n\n"
+                "## Reporting Vulnerabilities\n\n"
+                "Please report security vulnerabilities to: security@icdev.ai\n\n"
+                "## Supported Versions\n\n"
+                "| Version | Supported |\n"
+                "| ------- | --------- |\n"
+                "| latest  | Yes       |\n",
+                encoding="utf-8",
+            )
+        changes.append(change)
+
+    return changes
+
+
+def _remediate_least_privilege(
+    target_path: str, project_id: Optional[str], dry_run: bool
+) -> List[Dict]:
+    """Add security hardening to Dockerfiles and compose files."""
+    changes = []
+
+    # Find Dockerfiles
+    for fname in ("Dockerfile", "docker-compose.yml", "docker-compose.yaml"):
+        fpath = Path(target_path) / fname
+        if not fpath.exists():
+            continue
+
+        content = fpath.read_text(encoding="utf-8")
+
+        if fname == "Dockerfile" and "USER " not in content:
+            change = {
+                "file": fname,
+                "action": "add_non_root_user",
+                "applied": not dry_run,
+            }
+            if not dry_run:
+                # Add non-root user before CMD/ENTRYPOINT
+                lines = content.split("\n")
+                insert_idx = len(lines) - 1
+                for i, line in enumerate(lines):
+                    if line.strip().startswith(("CMD", "ENTRYPOINT")):
+                        insert_idx = i
+                        break
+                user_lines = [
+                    "",
+                    "# Security: run as non-root",
+                    "RUN addgroup --system appgroup && adduser --system --ingroup appgroup appuser",
+                    "USER appuser",
+                    "",
+                ]
+                for j, ul in enumerate(user_lines):
+                    lines.insert(insert_idx + j, ul)
+                fpath.write_text("\n".join(lines), encoding="utf-8")
+            changes.append(change)
+
+        if "compose" in fname and "read_only: true" not in content:
+            change = {
+                "file": fname,
+                "action": "add_read_only_fs_note",
+                "applied": not dry_run,
+                "note": "Manual review recommended — add read_only: true to service volumes",
+            }
+            changes.append(change)
+
+    return changes
+
+
+def _remediate_input_validation(
+    target_path: str, project_id: Optional[str], dry_run: bool
+) -> List[Dict]:
+    """Flag files with unsafe input patterns for review."""
+    changes = []
+    dangerous = [
+        (r"\beval\s*\(", "eval()"),
+        (r"\bexec\s*\(", "exec()"),
+        (r"os\.system\s*\(", "os.system()"),
+    ]
+    skip_dirs = {
+        ".git", "node_modules", "venv", "__pycache__", ".tmp", "data",
+        ".venv", "env", "dist", "build", "site-packages",
+    }
+
+    for root, dirs, files in os.walk(target_path):
+        dirs[:] = [d for d in dirs if d not in skip_dirs]
+        for fname in files:
+            if not fname.endswith(".py"):
+                continue
+            fpath = os.path.join(root, fname)
+            try:
+                with open(fpath, "r", encoding="utf-8", errors="ignore") as f:
+                    content = f.read(50000)
+                for pattern, label in dangerous:
+                    if re.search(pattern, content):
+                        rel_path = os.path.relpath(fpath, target_path)
+                        changes.append({
+                            "file": rel_path,
+                            "action": "flag_for_review",
+                            "pattern": label,
+                            "applied": False,
+                            "note": f"Contains {label} — manual review required",
+                        })
+            except (OSError, UnicodeDecodeError):
+                continue
+
+    return changes
+
+
+def _remediate_secrets(
+    target_path: str, project_id: Optional[str], dry_run: bool
+) -> List[Dict]:
+    """Flag files containing hardcoded secrets (manual remediation only)."""
+    # Secrets remediation is always manual — we only flag, never auto-modify
+    return [{
+        "action": "manual_review_required",
+        "applied": False,
+        "note": "Secret rotation requires manual review. Move secrets to environment variables or a secrets manager.",
+    }]
+
+
+def _remediate_sbom(
+    target_path: str, project_id: Optional[str], dry_run: bool
+) -> List[Dict]:
+    """Generate SBOM for the project."""
+    changes = []
+    if dry_run:
+        return [{"action": "generate_sbom", "applied": False, "note": "Would generate CycloneDX SBOM"}]
+
+    try:
+        import subprocess, sys
+        result = subprocess.run(
+            [sys.executable, "-c",
+             f"import sys; sys.path.insert(0,'.'); "
+             f"from tools.compliance.sbom_generator import generate_sbom; "
+             f"import json; r = generate_sbom(project_id='{project_id or 'audit-heal'}'); "
+             f"print(json.dumps({{'count': len(r.get('components', []))}}))"],
+            capture_output=True, text=True, timeout=15,
+            cwd=str(BASE_DIR),
+            env={**os.environ, "ICDEV_STORAGE_BACKEND": "sqlite"},
+        )
+        if result.returncode == 0:
+            data = json.loads(result.stdout.strip())
+            changes.append({
+                "action": "generate_sbom",
+                "applied": True,
+                "components": data.get("count", 0),
+            })
+        else:
+            changes.append({"action": "generate_sbom", "applied": False, "note": "SBOM tool unavailable"})
+    except Exception as e:
+        changes.append({"action": "generate_sbom", "applied": False, "note": str(e)})
+
+    return changes
+
+
+def _remediate_vuln_scanning(
+    target_path: str, project_id: Optional[str], dry_run: bool
+) -> List[Dict]:
+    """Create CI/CD scanning configuration."""
+    changes = []
+
+    # Create a basic scanning config
+    scan_config_path = Path(target_path) / ".security-scan.yml"
+    if not scan_config_path.exists():
+        change = {"file": ".security-scan.yml", "action": "create", "applied": not dry_run}
+        if not dry_run:
+            scan_config_path.write_text(
+                "# Security Scanning Configuration\n"
+                "scanners:\n"
+                "  sast: true\n"
+                "  dependency_audit: true\n"
+                "  secret_detection: true\n"
+                "  container_scan: false\n"
+                "\n"
+                "# Run on every commit\n"
+                "triggers:\n"
+                "  - push\n"
+                "  - pull_request\n",
+                encoding="utf-8",
+            )
+        changes.append(change)
+
+    return changes
+
+
+def _remediate_dependencies(
+    target_path: str, project_id: Optional[str], dry_run: bool
+) -> List[Dict]:
+    """Flag outdated dependencies for update (manual action)."""
+    return [{
+        "action": "manual_review_required",
+        "applied": False,
+        "note": "Run 'pip list --outdated' and update vulnerable packages. Review SBOM for CVE details.",
+    }]
+
+
+def _remediate_encryption(
+    target_path: str, project_id: Optional[str], dry_run: bool
+) -> List[Dict]:
+    """Add encryption configuration guidance."""
+    return [{
+        "action": "manual_review_required",
+        "applied": False,
+        "note": "Enable TLS for all network communication. Add SSL context to HTTP servers and enforce HTTPS redirects.",
+    }]
+
+
+# ---------------------------------------------------------------
+# Exclusions (Risk Acceptance)
+# ---------------------------------------------------------------
+
+def create_exclusion(
+    control_id: str,
+    reason: str,
+    owner: str,
+    project_id: Optional[str] = None,
+    duration_days: int = 90,
+    risk_level: str = "accepted",
+) -> Dict:
+    """Register a risk acceptance exclusion for a control.
+
+    Excluded controls are removed from scoring and flagged in reports.
+    Exclusions expire after duration_days and must be renewed.
+
+    Args:
+        control_id: Control ID to exclude (e.g., 'AC-6', 'SI-10').
+        reason: Business justification for risk acceptance.
+        owner: Person/role responsible for the risk.
+        project_id: Project scope for the exclusion.
+        duration_days: Exclusion validity period (max 365).
+        risk_level: One of 'accepted', 'mitigated', 'transferred'.
+    """
+    timestamp = datetime.now(timezone.utc)
+    duration_days = min(duration_days, 365)  # Cap at 1 year
+    expiry = datetime.fromtimestamp(
+        timestamp.timestamp() + duration_days * 86400, tz=timezone.utc
+    ).isoformat()
+
+    exclusion = {
+        "control_id": control_id,
+        "reason": reason,
+        "owner": owner,
+        "project_id": project_id,
+        "risk_level": risk_level,
+        "created": timestamp.isoformat(),
+        "expires": expiry,
+        "duration_days": duration_days,
+        "status": "active",
+    }
+
+    if get_connection:
+        try:
+            with get_connection() as conn:
+                conn.execute(
+                    """CREATE TABLE IF NOT EXISTS audit_engine_exclusions (
+                        id INTEGER PRIMARY KEY AUTOINCREMENT,
+                        control_id TEXT NOT NULL,
+                        reason TEXT NOT NULL,
+                        owner TEXT NOT NULL,
+                        project_id TEXT,
+                        risk_level TEXT DEFAULT 'accepted',
+                        created TEXT NOT NULL,
+                        expires TEXT NOT NULL,
+                        duration_days INTEGER,
+                        status TEXT DEFAULT 'active',
+                        classification TEXT DEFAULT 'CUI'
+                    )""",
+                )
+                conn.execute(
+                    """INSERT INTO audit_engine_exclusions
+                       (control_id, reason, owner, project_id, risk_level,
+                        created, expires, duration_days, status)
+                       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+                    (
+                        control_id, reason, owner, project_id, risk_level,
+                        timestamp.isoformat(), expiry, duration_days, "active",
+                    ),
+                )
+                # Also log to audit trail
+                conn.execute(
+                    """INSERT INTO audit_trail (event_type, event_data, timestamp, classification)
+                       VALUES (?, ?, ?, ?)""",
+                    (
+                        "audit_engine_exclusion_created",
+                        json.dumps(exclusion),
+                        timestamp.isoformat(),
+                        "CUI",
+                    ),
+                )
+            exclusion["stored"] = True
+        except Exception as e:
+            exclusion["stored"] = False
+            exclusion["storage_error"] = str(e)
+    else:
+        exclusion["stored"] = False
+
+    return exclusion
+
+
+def list_exclusions(project_id: Optional[str] = None) -> List[Dict]:
+    """List active and expired exclusions."""
+    if not get_connection:
+        return []
+    try:
+        with get_connection() as conn:
+            conn.execute(
+                """CREATE TABLE IF NOT EXISTS audit_engine_exclusions (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    control_id TEXT NOT NULL,
+                    reason TEXT NOT NULL,
+                    owner TEXT NOT NULL,
+                    project_id TEXT,
+                    risk_level TEXT DEFAULT 'accepted',
+                    created TEXT NOT NULL,
+                    expires TEXT NOT NULL,
+                    duration_days INTEGER,
+                    status TEXT DEFAULT 'active',
+                    classification TEXT DEFAULT 'CUI'
+                )""",
+            )
+            if project_id:
+                rows = conn.execute(
+                    "SELECT * FROM audit_engine_exclusions WHERE project_id = ? ORDER BY created DESC",
+                    (project_id,),
+                ).fetchall()
+            else:
+                rows = conn.execute(
+                    "SELECT * FROM audit_engine_exclusions ORDER BY created DESC"
+                ).fetchall()
+
+        now = datetime.now(timezone.utc).isoformat()
+        result = []
+        for row in rows:
+            excl = dict(row) if hasattr(row, "keys") else {
+                "id": row[0], "control_id": row[1], "reason": row[2],
+                "owner": row[3], "project_id": row[4], "risk_level": row[5],
+                "created": row[6], "expires": row[7], "duration_days": row[8],
+                "status": row[9],
+            }
+            # Check if expired
+            if excl.get("expires", "") < now and excl.get("status") == "active":
+                excl["status"] = "expired"
+            result.append(excl)
+        return result
+    except Exception:
+        return []
+
+
+def get_excluded_controls(project_id: Optional[str] = None) -> set:
+    """Get set of currently excluded (active, non-expired) control IDs."""
+    exclusions = list_exclusions(project_id)
+    now = datetime.now(timezone.utc).isoformat()
+    return {
+        e["control_id"]
+        for e in exclusions
+        if e.get("status") == "active" and e.get("expires", "") >= now
+    }
+
+
+# ---------------------------------------------------------------
+# POAM Integration
+# ---------------------------------------------------------------
+
+def generate_poam_entries(
+    audit_result: Dict,
+    project_id: Optional[str] = None,
+) -> Dict:
+    """Generate POAM entries from audit findings.
+
+    Creates Plan of Action & Milestones entries for all failing controls,
+    with severity-based timelines and cross-regime impact analysis.
+
+    Returns:
+        Dict with poam_entries list and summary.
+    """
+    entries = []
+    excluded = get_excluded_controls(project_id)
+
+    # Collect all failing controls across regimes with dedup
+    seen_controls = {}
+    for regime_id, regime_data in audit_result.get("regime_scores", {}).items():
+        if regime_data.get("error"):
+            continue
+        for cat in regime_data.get("categories", []):
+            for ctrl in cat.get("controls", []):
+                ctrl_id = ctrl.get("control_id", "")
+                status = ctrl.get("status", "")
+                if status not in ("fail", "partial"):
+                    continue
+                if ctrl_id in excluded:
+                    continue
+
+                if ctrl_id not in seen_controls:
+                    seen_controls[ctrl_id] = {
+                        "control_id": ctrl_id,
+                        "control_name": ctrl.get("control_name", ctrl_id),
+                        "status": status,
+                        "severity": ctrl.get("severity", "medium"),
+                        "regimes": [],
+                        "categories": [],
+                    }
+                seen_controls[ctrl_id]["regimes"].append(regime_id)
+                if cat.get("category_name") not in seen_controls[ctrl_id]["categories"]:
+                    seen_controls[ctrl_id]["categories"].append(cat.get("category_name", ""))
+
+    # Build POAM entries with timelines
+    severity_timelines = {
+        "critical": {"days": 30, "priority": "P1"},
+        "high": {"days": 90, "priority": "P2"},
+        "medium": {"days": 180, "priority": "P3"},
+        "low": {"days": 365, "priority": "P4"},
+    }
+
+    now = datetime.now(timezone.utc)
+    for ctrl_id, ctrl_data in seen_controls.items():
+        sev = ctrl_data["severity"]
+        timeline = severity_timelines.get(sev, severity_timelines["medium"])
+
+        due_date = datetime.fromtimestamp(
+            now.timestamp() + timeline["days"] * 86400, tz=timezone.utc
+        ).strftime("%Y-%m-%d")
+
+        # Get recommendation if available
+        rec = REMEDIATION_ACTIONS.get(
+            _extract_control_family(ctrl_id) or "",
+            {}
+        )
+
+        entry = {
+            "poam_id": f"POAM-AE-{ctrl_id.replace('.', '-')}",
+            "control_id": ctrl_id,
+            "control_name": ctrl_data["control_name"],
+            "weakness_description": f"Control {ctrl_id} is {'not implemented' if ctrl_data['status'] == 'fail' else 'partially implemented'}",
+            "severity": sev,
+            "priority": timeline["priority"],
+            "status": "Open",
+            "affected_regimes": ctrl_data["regimes"],
+            "regime_count": len(ctrl_data["regimes"]),
+            "scheduled_completion": due_date,
+            "timeline_days": timeline["days"],
+            "milestones": _generate_milestones(ctrl_id, rec, timeline["days"]),
+            "remediation_type": rec.get("action_type", "manual"),
+            "auto_fixable": ctrl_id in {
+                _extract_control_family(k) or k
+                for k in REMEDIATION_ACTIONS
+            } if _extract_control_family(ctrl_id) else False,
+            "resources_required": rec.get("description", "Manual assessment required"),
+        }
+        entries.append(entry)
+
+    # Sort by priority then regime count
+    entries.sort(key=lambda e: (e["priority"], -e["regime_count"]))
+
+    # Write POAM to DB if possible
+    poam_stored = False
+    if get_connection and project_id:
+        try:
+            with get_connection() as conn:
+                for entry in entries:
+                    conn.execute(
+                        """INSERT INTO audit_trail (event_type, event_data, timestamp, classification)
+                           VALUES (?, ?, ?, ?)""",
+                        (
+                            "audit_engine_poam_generated",
+                            json.dumps(entry),
+                            now.isoformat(),
+                            "CUI",
+                        ),
+                    )
+            poam_stored = True
+        except Exception:
+            pass
+
+    return {
+        "success": True,
+        "project_id": project_id,
+        "timestamp": now.isoformat(),
+        "total_entries": len(entries),
+        "by_priority": {
+            p: len([e for e in entries if e["priority"] == p])
+            for p in ("P1", "P2", "P3", "P4")
+        },
+        "excluded_controls": len(excluded),
+        "poam_entries": entries,
+        "stored": poam_stored,
+    }
+
+
+def _generate_milestones(
+    ctrl_id: str, rec: Dict, total_days: int
+) -> List[Dict]:
+    """Generate milestone schedule for a POAM entry."""
+    milestones = [
+        {
+            "milestone": "Assessment & Planning",
+            "target_date_offset_days": max(7, total_days // 6),
+            "description": f"Assess current state of {ctrl_id} and develop remediation plan",
+        },
+        {
+            "milestone": "Implementation",
+            "target_date_offset_days": max(14, total_days // 2),
+            "description": rec.get("description", f"Implement remediation for {ctrl_id}"),
+        },
+        {
+            "milestone": "Verification & Testing",
+            "target_date_offset_days": max(21, total_days * 3 // 4),
+            "description": f"Verify {ctrl_id} implementation and run compliance check",
+        },
+        {
+            "milestone": "Closure",
+            "target_date_offset_days": total_days,
+            "description": f"Close POAM entry after successful re-audit of {ctrl_id}",
+        },
+    ]
+    return milestones
+
+
+# ---------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------
+
+def _extract_control_family(control_id: str) -> Optional[str]:
+    """Extract NIST control family from control ID.
+
+    Handles multiple ID formats:
+      - NIST 800-53: AC-2, AU-3, SI-10, AC-2(1)
+      - CMMC: AC.L2-3.1.1, SI.L2-3.14.1
+      - CISA SBD: SBD-MS-01, SBD-VT-02
+      - NIST 800-171: 3.1.1, 3.5.2
+      - DoD CSSP/other: CSSP-1.1, etc.
+    """
+    family_map = {
+        "AC": "AC-6", "AU": "AU-2", "CM": "CM-6", "IA": "IA-5",
+        "SI": "SI-10", "SC": "SC-13", "SA": "SA-11", "RA": "RA-5",
+        "SR": "SA-11", "AT": "AU-2", "MP": "SC-13", "PE": "AC-6",
+        "PL": "SA-11", "PS": "AC-6", "CA": "RA-5", "CP": "CM-6",
+        "IR": "SI-2", "MA": "CM-6", "PM": "RA-5",
+    }
+
+    # NIST 800-53 style: AC-2, AU-3, SI-10
+    m = re.match(r"^([A-Z]{2})-(\d+)", control_id)
+    if m:
+        return f"{m.group(1)}-{m.group(2)}"
+
+    # CMMC style: AC.L2-3.1.1, SI.L2-3.14.1
+    m = re.match(r"^([A-Z]{2})\.L\d+-", control_id)
+    if m:
+        return family_map.get(m.group(1))
+
+    # CISA SBD style: SBD-MS-01, SBD-VT-02, etc.
+    sbd_to_nist = {
+        "MS": "CM-6",   # Memory Safety -> Config Mgmt
+        "VT": "RA-5",   # Vulnerability Transparency -> Risk Assessment
+        "SD": "SA-11",   # Secure Development -> Dev Testing
+        "DF": "CM-6",   # Default Security -> Config Mgmt
+        "MF": "IA-5",   # Multi-Factor Auth -> Auth Mgmt
+        "IL": "AU-2",   # Intrusion Logging -> Audit Logging
+        "EP": "SI-2",   # Evidence Patching -> Flaw Remediation
+        "DP": "SC-8",   # Data Protection -> Comms Protection
+        "CV": "SA-11",  # CVD Policy -> Dev Testing
+    }
+    m = re.match(r"^SBD-([A-Z]{2})-\d+", control_id)
+    if m:
+        return sbd_to_nist.get(m.group(1), "SA-11")
+
+    # NIST 800-171 style: 3.1.1, 3.5.2 — map chapter to NIST family
+    nist171_chapter_map = {
+        "1": "AC-6",   # Access Control
+        "2": "AU-2",   # Awareness & Training
+        "3": "AU-2",   # Audit & Accountability
+        "4": "CM-6",   # Configuration Mgmt
+        "5": "IA-5",   # Identification & Auth
+        "6": "SI-2",   # Incident Response
+        "7": "CM-6",   # Maintenance
+        "8": "SC-13",  # Media Protection
+        "9": "AC-6",   # Physical Protection
+        "10": "AC-6",  # Personnel Security
+        "11": "RA-5",  # Risk Assessment
+        "12": "RA-5",  # Security Assessment
+        "13": "SC-13", # System & Comms Protection
+        "14": "SI-2",  # System & Info Integrity
+    }
+    m = re.match(r"^3\.(\d+)\.\d+", control_id)
+    if m:
+        return nist171_chapter_map.get(m.group(1), "SA-11")
+
+    # Try just the family prefix (e.g., XX-something)
+    m = re.match(r"^([A-Z]{2})-", control_id)
+    if m:
+        return family_map.get(m.group(1))
+
+    return None
+
+
+def _count_affected_regimes(audit_result: Dict, control_family: str) -> int:
+    """Count how many regimes have failing controls in this family."""
+    count = 0
+    prefix = control_family.split("-")[0] if "-" in control_family else control_family
+
+    for regime_data in audit_result.get("regime_scores", {}).values():
+        if regime_data.get("error"):
+            continue
+        for cat in regime_data.get("categories", []):
+            for ctrl in cat.get("controls", []):
+                ctrl_id = ctrl.get("control_id", "")
+                if ctrl_id.startswith(prefix) and ctrl.get("status") in ("fail", "partial"):
+                    count += 1
+                    break
+    return count
+
+
+def _log_remediation(result: Dict):
+    """Log remediation action to audit trail (append-only)."""
+    if not get_connection:
+        return
+    try:
+        with get_connection() as conn:
+            conn.execute(
+                """INSERT INTO audit_trail (event_type, event_data, timestamp, classification)
+                   VALUES (?, ?, ?, ?)""",
+                (
+                    "audit_engine_self_heal",
+                    json.dumps({
+                        "control_id": result["control_id"],
+                        "action_title": result["action_title"],
+                        "files_modified": result.get("files_modified", 0),
+                        "target_path": result["target_path"],
+                        "dry_run": result["dry_run"],
+                    }),
+                    result["timestamp"],
+                    "CUI",
+                ),
+            )
+    except Exception:
+        pass  # Audit trail table may not exist for BYOS
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Audit Engine Self-Heal")
+    parser.add_argument("--control", required=True, help="Control ID to remediate")
+    parser.add_argument("--target", default=".", help="Target directory")
+    parser.add_argument("--project-id", help="Project ID")
+    parser.add_argument("--dry-run", action="store_true", help="Preview without applying")
+    parser.add_argument("--list", action="store_true", help="List available actions")
+    parser.add_argument("--json", action="store_true", help="JSON output")
+    args = parser.parse_args()
+
+    if args.list:
+        print(json.dumps(list(REMEDIATION_ACTIONS.values()), indent=2))
+    else:
+        result = execute_remediation(args.control, args.target, args.project_id, args.dry_run)
+        if args.json:
+            print(json.dumps(result, indent=2))
+        else:
+            status = "SUCCESS" if result["success"] else "FAILED"
+            print(f"\n{status}: {result.get('action_title', args.control)}")
+            if result.get("error"):
+                print(f"  Error: {result['error']}")
+            for c in result.get("changes", []):
+                applied = "APPLIED" if c.get("applied") else "PREVIEW"
+                print(f"  [{applied}] {c.get('file', '')} — {c.get('action', '')}")
+                if c.get("note"):
+                    print(f"           {c['note']}")