icdev 0.0.3__py3-none-any.whl

tools/compliance/narrative_workflow.py

@@ -0,0 +1,814 @@
+#!/usr/bin/env python3
+# CUI // SP-CTI
+# ////////////////////////////////////////////////////////////////////
+# CONTROLLED UNCLASSIFIED INFORMATION (CUI) // SP-CTI
+# Distribution: Distribution D -- Authorized DoD Personnel Only
+# ////////////////////////////////////////////////////////////////////
+"""AI-Generated Compliance Narratives — approval workflow engine.
+
+Manages the lifecycle of generated SSP/POAM narratives through a
+draft -> pending_review -> approved/rejected workflow. Supports
+template-based generation (deterministic) and optional LLM-enhanced
+generation via the two-tier router.
+
+Database tables:
+  - narrative_approvals (read/write) — narrative lifecycle records
+  - project_controls (read) — mapped controls per project
+  - compliance_controls (read) — control definitions
+"""
+
+import argparse
+import json
+import sqlite3
+import sys
+import uuid
+from datetime import datetime, timezone
+from pathlib import Path
+
+BASE_DIR = Path(__file__).resolve().parent.parent.parent
+DB_PATH = BASE_DIR / "data" / "icdev.db"
+
+try:
+    from tools.compat.db_utils import get_db_connection
+except ImportError:
+    get_db_connection = None
+
+
+def _get_connection(db_path=None):
+    """Get a database connection with row factory."""
+    if get_db_connection:
+        return get_db_connection(db_path or DB_PATH, validate=True)
+    path = db_path or DB_PATH
+    if not Path(path).exists():
+        raise FileNotFoundError(
+            f"Database not found: {path}\n"
+            "Run: python tools/db/init_icdev_db.py"
+        )
+    conn = sqlite3.connect(str(path))
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+def _ensure_tables(conn):
+    """Ensure the narrative_approvals table exists."""
+    conn.execute("""
+        CREATE TABLE IF NOT EXISTS narrative_approvals (
+            id TEXT PRIMARY KEY,
+            project_id TEXT NOT NULL,
+            control_id TEXT NOT NULL,
+            narrative_type TEXT NOT NULL DEFAULT 'control',
+            template_draft TEXT,
+            llm_draft TEXT,
+            llm_reviewed TEXT,
+            final_narrative TEXT,
+            status TEXT NOT NULL DEFAULT 'draft',
+            reviewer TEXT,
+            comment TEXT,
+            created_at TEXT NOT NULL,
+            updated_at TEXT NOT NULL
+        )
+    """)
+    conn.commit()
+
+
+def _log_audit_event(conn, project_id, event_type, details):
+    """Log an audit trail event (append-only, NIST AU compliant)."""
+    try:
+        conn.execute(
+            """INSERT INTO audit_trail
+               (project_id, event_type, actor, action, details, classification)
+               VALUES (?, ?, ?, ?, ?, ?)""",
+            (
+                project_id,
+                event_type,
+                "icdev-narrative-workflow",
+                event_type,
+                json.dumps(details, default=str),
+                "CUI",
+            ),
+        )
+        conn.commit()
+    except Exception as e:
+        print(f"Warning: Could not log audit event: {e}", file=sys.stderr)
+
+
+def _get_control_info(conn, project_id, control_id):
+    """Fetch control metadata from project_controls and compliance_controls.
+
+    Returns:
+        Dict with control_id, title, description, implementation_status,
+        and implementation_description. Returns minimal dict if not found.
+    """
+    info = {
+        "control_id": control_id,
+        "title": control_id,
+        "description": "",
+        "implementation_status": "Not Assessed",
+        "implementation_description": "",
+    }
+
+    # Try project_controls
+    try:
+        row = conn.execute(
+            """SELECT control_id, implementation_status, implementation_description
+               FROM project_controls
+               WHERE project_id = ? AND control_id = ?
+               LIMIT 1""",
+            (project_id, control_id),
+        ).fetchone()
+        if row:
+            info["implementation_status"] = row["implementation_status"] or "Not Assessed"
+            info["implementation_description"] = row["implementation_description"] or ""
+    except sqlite3.OperationalError:
+        pass
+
+    # Try compliance_controls for title/description
+    try:
+        row = conn.execute(
+            """SELECT title, description FROM compliance_controls
+               WHERE control_id = ?
+               LIMIT 1""",
+            (control_id,),
+        ).fetchone()
+        if row:
+            info["title"] = row["title"] or control_id
+            info["description"] = row["description"] or ""
+    except sqlite3.OperationalError:
+        pass
+
+    return info
+
+
+def _generate_template_draft(control_info):
+    """Generate a deterministic template-based narrative draft.
+
+    Args:
+        control_info: Dict from _get_control_info.
+
+    Returns:
+        String narrative draft.
+    """
+    # Try importing NarrativeGenerator from narrative_generator
+    try:
+        from tools.compliance.narrative_generator import NarrativeGenerator
+        gen = NarrativeGenerator()
+        result = gen.generate(
+            control_id=control_info["control_id"],
+            context=control_info,
+        )
+        if isinstance(result, dict) and result.get("narrative"):
+            return result["narrative"]
+        if isinstance(result, str) and result.strip():
+            return result
+    except (ImportError, AttributeError, TypeError, Exception):
+        pass
+
+    # Fallback: deterministic template
+    lines = [
+        f"## {control_info['control_id']} — {control_info['title']}",
+        "",
+        f"**Implementation Status:** {control_info['implementation_status']}",
+        "",
+    ]
+    if control_info["description"]:
+        lines.append(f"**Control Description:** {control_info['description']}")
+        lines.append("")
+    if control_info["implementation_description"]:
+        lines.append("### Implementation Narrative")
+        lines.append("")
+        lines.append(control_info["implementation_description"])
+        lines.append("")
+    else:
+        lines.append("### Implementation Narrative")
+        lines.append("")
+        lines.append(
+            f"The organization implements {control_info['control_id']} "
+            f"({control_info['title']}) through a combination of technical "
+            "controls and organizational policies. Detailed implementation "
+            "evidence is maintained in the project compliance repository."
+        )
+        lines.append("")
+
+    return "\n".join(lines)
+
+
+def _generate_llm_draft(control_info, template_draft):
+    """Generate an LLM-enhanced narrative draft using two-tier routing.
+
+    Args:
+        control_info: Dict from _get_control_info.
+        template_draft: The template-based draft to enhance.
+
+    Returns:
+        String LLM-enhanced draft, or None if LLM not available.
+    """
+    try:
+        from tools.llm.router import LLMRouter
+        router = LLMRouter()
+        prompt = (
+            f"You are a compliance narrative writer for NIST 800-53. "
+            f"Enhance the following SSP control narrative for "
+            f"{control_info['control_id']} ({control_info['title']}). "
+            f"Make it suitable for a System Security Plan at IL4. "
+            f"Maintain factual accuracy and keep all technical details.\n\n"
+            f"Draft to enhance:\n{template_draft}"
+        )
+        response = router.invoke(
+            function_name="compliance_export",
+            prompt=prompt,
+        )
+        if isinstance(response, dict):
+            return response.get("text") or response.get("content") or str(response)
+        if isinstance(response, str) and response.strip():
+            return response
+    except (ImportError, AttributeError, TypeError, Exception):
+        pass
+
+    return None
+
+
+# --------------------------------------------------------------------------
+# Public API functions
+# --------------------------------------------------------------------------
+
+
+def generate_narrative(project_id, control_id, narrative_type="control",
+                       use_llm=False, db_path=None):
+    """Generate a compliance narrative for a control.
+
+    Creates a template-based draft and optionally an LLM-enhanced draft.
+    INSERTs into narrative_approvals with status='draft'.
+
+    Args:
+        project_id: Project identifier.
+        control_id: NIST 800-53 control ID (e.g. 'AC-2').
+        narrative_type: Type of narrative ('control', 'ssp_section', 'poam').
+        use_llm: If True, also generate an LLM-enhanced draft.
+        db_path: Optional database path override.
+
+    Returns:
+        Dict with id, project_id, control_id, narrative_type, template_draft,
+        llm_draft, and status.
+    """
+    conn = _get_connection(db_path)
+    try:
+        _ensure_tables(conn)
+
+        control_info = _get_control_info(conn, project_id, control_id)
+        template_draft = _generate_template_draft(control_info)
+
+        llm_draft = None
+        if use_llm:
+            llm_draft = _generate_llm_draft(control_info, template_draft)
+
+        now = datetime.now(timezone.utc).isoformat()
+        narrative_id = f"nar-{uuid.uuid4().hex[:12]}"
+
+        conn.execute(
+            """INSERT INTO narrative_approvals
+               (id, project_id, control_id, narrative_type, template_draft,
+                llm_draft, status, created_at, updated_at)
+               VALUES (?, ?, ?, ?, ?, ?, 'draft', ?, ?)""",
+            (
+                narrative_id,
+                project_id,
+                control_id,
+                narrative_type,
+                template_draft,
+                llm_draft,
+                now,
+                now,
+            ),
+        )
+        conn.commit()
+
+        # Audit trail
+        _log_audit_event(conn, project_id, "narrative_generated", {
+            "narrative_id": narrative_id,
+            "control_id": control_id,
+            "narrative_type": narrative_type,
+            "used_llm": use_llm,
+        })
+
+        result = {
+            "id": narrative_id,
+            "project_id": project_id,
+            "control_id": control_id,
+            "narrative_type": narrative_type,
+            "template_draft": template_draft,
+            "llm_draft": llm_draft,
+            "status": "draft",
+        }
+
+        print(f"Narrative generated: {narrative_id} for {control_id} "
+              f"(type={narrative_type}, llm={use_llm})")
+
+        return result
+
+    finally:
+        conn.close()
+
+
+def submit_for_review(narrative_id, db_path=None):
+    """Submit a draft narrative for review.
+
+    Updates narrative_approvals SET status='pending_review'.
+
+    Args:
+        narrative_id: The narrative record ID.
+        db_path: Optional database path override.
+
+    Returns:
+        Dict with id and status.
+    """
+    conn = _get_connection(db_path)
+    try:
+        _ensure_tables(conn)
+
+        row = conn.execute(
+            "SELECT id, project_id, control_id, status FROM narrative_approvals WHERE id = ?",
+            (narrative_id,),
+        ).fetchone()
+        if not row:
+            raise ValueError(f"Narrative '{narrative_id}' not found.")
+        if row["status"] != "draft":
+            raise ValueError(
+                f"Narrative '{narrative_id}' is '{row['status']}', "
+                "only 'draft' can be submitted for review."
+            )
+
+        now = datetime.now(timezone.utc).isoformat()
+        conn.execute(
+            """UPDATE narrative_approvals
+               SET status = 'pending_review', updated_at = ?
+               WHERE id = ?""",
+            (now, narrative_id),
+        )
+        conn.commit()
+
+        # Audit trail
+        _log_audit_event(conn, row["project_id"], "narrative_submitted", {
+            "narrative_id": narrative_id,
+            "control_id": row["control_id"],
+        })
+
+        result = {
+            "id": narrative_id,
+            "status": "pending_review",
+        }
+
+        print(f"Narrative submitted for review: {narrative_id}")
+
+        return result
+
+    finally:
+        conn.close()
+
+
+def approve_narrative(narrative_id, reviewer, comment=None, db_path=None):
+    """Approve a narrative pending review.
+
+    Sets status='approved' and populates final_narrative with the
+    llm_reviewed draft if available, otherwise the template_draft.
+
+    Args:
+        narrative_id: The narrative record ID.
+        reviewer: Name/ID of the reviewer.
+        comment: Optional review comment.
+        db_path: Optional database path override.
+
+    Returns:
+        Dict with id, status, and reviewer.
+    """
+    conn = _get_connection(db_path)
+    try:
+        _ensure_tables(conn)
+
+        row = conn.execute(
+            """SELECT id, project_id, control_id, status, template_draft,
+                      llm_draft, llm_reviewed
+               FROM narrative_approvals WHERE id = ?""",
+            (narrative_id,),
+        ).fetchone()
+        if not row:
+            raise ValueError(f"Narrative '{narrative_id}' not found.")
+        if row["status"] != "pending_review":
+            raise ValueError(
+                f"Narrative '{narrative_id}' is '{row['status']}', "
+                "only 'pending_review' can be approved."
+            )
+
+        # Pick the best available narrative as final
+        final = (
+            row["llm_reviewed"]
+            or row["llm_draft"]
+            or row["template_draft"]
+            or ""
+        )
+
+        now = datetime.now(timezone.utc).isoformat()
+        conn.execute(
+            """UPDATE narrative_approvals
+               SET status = 'approved', final_narrative = ?,
+                   reviewer = ?, comment = ?, updated_at = ?
+               WHERE id = ?""",
+            (final, reviewer, comment, now, narrative_id),
+        )
+        conn.commit()
+
+        # Audit trail
+        _log_audit_event(conn, row["project_id"], "narrative_approved", {
+            "narrative_id": narrative_id,
+            "control_id": row["control_id"],
+            "reviewer": reviewer,
+            "comment": comment,
+        })
+
+        result = {
+            "id": narrative_id,
+            "status": "approved",
+            "reviewer": reviewer,
+        }
+
+        print(f"Narrative approved: {narrative_id} by {reviewer}")
+
+        return result
+
+    finally:
+        conn.close()
+
+
+def reject_narrative(narrative_id, reviewer, comment, db_path=None):
+    """Reject a narrative pending review.
+
+    Sets status='rejected' with the reviewer's comment.
+
+    Args:
+        narrative_id: The narrative record ID.
+        reviewer: Name/ID of the reviewer.
+        comment: Rejection reason (required).
+        db_path: Optional database path override.
+
+    Returns:
+        Dict with id, status, reviewer, and comment.
+    """
+    conn = _get_connection(db_path)
+    try:
+        _ensure_tables(conn)
+
+        row = conn.execute(
+            """SELECT id, project_id, control_id, status
+               FROM narrative_approvals WHERE id = ?""",
+            (narrative_id,),
+        ).fetchone()
+        if not row:
+            raise ValueError(f"Narrative '{narrative_id}' not found.")
+        if row["status"] != "pending_review":
+            raise ValueError(
+                f"Narrative '{narrative_id}' is '{row['status']}', "
+                "only 'pending_review' can be rejected."
+            )
+
+        now = datetime.now(timezone.utc).isoformat()
+        conn.execute(
+            """UPDATE narrative_approvals
+               SET status = 'rejected', reviewer = ?, comment = ?,
+                   updated_at = ?
+               WHERE id = ?""",
+            (reviewer, comment, now, narrative_id),
+        )
+        conn.commit()
+
+        # Audit trail
+        _log_audit_event(conn, row["project_id"], "narrative_rejected", {
+            "narrative_id": narrative_id,
+            "control_id": row["control_id"],
+            "reviewer": reviewer,
+            "comment": comment,
+        })
+
+        result = {
+            "id": narrative_id,
+            "status": "rejected",
+            "reviewer": reviewer,
+            "comment": comment,
+        }
+
+        print(f"Narrative rejected: {narrative_id} by {reviewer} — {comment}")
+
+        return result
+
+    finally:
+        conn.close()
+
+
+def list_pending(project_id, db_path=None):
+    """List all narratives pending review for a project.
+
+    Args:
+        project_id: Project identifier.
+        db_path: Optional database path override.
+
+    Returns:
+        Dict with project_id, pending list, and count.
+    """
+    conn = _get_connection(db_path)
+    try:
+        _ensure_tables(conn)
+
+        rows = conn.execute(
+            """SELECT id, project_id, control_id, narrative_type,
+                      created_at, updated_at
+               FROM narrative_approvals
+               WHERE project_id = ? AND status = 'pending_review'
+               ORDER BY updated_at DESC""",
+            (project_id,),
+        ).fetchall()
+
+        pending = []
+        for row in rows:
+            pending.append({
+                "id": row["id"],
+                "project_id": row["project_id"],
+                "control_id": row["control_id"],
+                "narrative_type": row["narrative_type"],
+                "created_at": row["created_at"],
+                "updated_at": row["updated_at"],
+            })
+
+        result = {
+            "project_id": project_id,
+            "pending": pending,
+            "count": len(pending),
+        }
+
+        print(f"Pending narratives for {project_id}: {len(pending)}")
+
+        return result
+
+    finally:
+        conn.close()
+
+
+def batch_generate(project_id, use_llm=False, db_path=None):
+    """Generate narratives for all controls without an approved narrative.
+
+    Queries project_controls and skips any control that already has an
+    'approved' entry in narrative_approvals.
+
+    Args:
+        project_id: Project identifier.
+        use_llm: If True, generate LLM-enhanced drafts.
+        db_path: Optional database path override.
+
+    Returns:
+        Dict with project_id, generated count, skipped count, and results.
+    """
+    conn = _get_connection(db_path)
+    try:
+        _ensure_tables(conn)
+
+        # Get all controls
+        ctrl_rows = conn.execute(
+            """SELECT DISTINCT control_id FROM project_controls
+               WHERE project_id = ?
+               ORDER BY control_id""",
+            (project_id,),
+        ).fetchall()
+
+        # Get already-approved controls
+        approved_rows = conn.execute(
+            """SELECT DISTINCT control_id FROM narrative_approvals
+               WHERE project_id = ? AND status = 'approved'""",
+            (project_id,),
+        ).fetchall()
+        approved_set = {r["control_id"] for r in approved_rows}
+
+        generated = 0
+        skipped = 0
+        results = []
+
+        for ctrl_row in ctrl_rows:
+            control_id = ctrl_row["control_id"]
+
+            if control_id in approved_set:
+                skipped += 1
+                continue
+
+            # Close connection before calling generate_narrative (it opens its own)
+            conn_db_path = db_path
+            try:
+                result = generate_narrative(
+                    project_id=project_id,
+                    control_id=control_id,
+                    narrative_type="control",
+                    use_llm=use_llm,
+                    db_path=conn_db_path,
+                )
+                results.append(result)
+                generated += 1
+            except Exception as e:
+                results.append({
+                    "control_id": control_id,
+                    "status": "error",
+                    "error": str(e),
+                })
+
+        result = {
+            "project_id": project_id,
+            "generated": generated,
+            "skipped": skipped,
+            "results": results,
+        }
+
+        print(f"Batch generate for {project_id}: "
+              f"{generated} generated, {skipped} skipped (already approved)")
+
+        return result
+
+    finally:
+        conn.close()
+
+
+# --------------------------------------------------------------------------
+# CLI entry point
+# --------------------------------------------------------------------------
+
+
+def main():
+    """CLI entry point for narrative approval workflow engine."""
+    parser = argparse.ArgumentParser(
+        description="AI-Generated Compliance Narratives — approval workflow engine"
+    )
+    parser.add_argument(
+        "--project-id", required=True,
+        help="Project ID in ICDEV database"
+    )
+    parser.add_argument(
+        "--db-path", type=Path, default=None,
+        help="Override database path"
+    )
+
+    # Action flags
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        "--generate", action="store_true",
+        help="Generate a narrative for a control"
+    )
+    group.add_argument(
+        "--submit", action="store_true",
+        help="Submit a narrative for review"
+    )
+    group.add_argument(
+        "--approve", action="store_true",
+        help="Approve a narrative"
+    )
+    group.add_argument(
+        "--reject", action="store_true",
+        help="Reject a narrative"
+    )
+    group.add_argument(
+        "--pending", action="store_true",
+        help="List all narratives pending review"
+    )
+    group.add_argument(
+        "--batch", action="store_true",
+        help="Batch generate narratives for all controls without approved narrative"
+    )
+
+    # Parameters for specific actions
+    parser.add_argument(
+        "--control-id", type=str, default=None,
+        help="Control ID for --generate"
+    )
+    parser.add_argument(
+        "--narrative-type", type=str, default="control",
+        help="Narrative type for --generate (default: control)"
+    )
+    parser.add_argument(
+        "--use-llm", action="store_true",
+        help="Use LLM for narrative generation (--generate/--batch)"
+    )
+    parser.add_argument(
+        "--narrative-id", type=str, default=None,
+        help="Narrative ID for --submit/--approve/--reject"
+    )
+    parser.add_argument(
+        "--reviewer", type=str, default=None,
+        help="Reviewer name for --approve/--reject"
+    )
+    parser.add_argument(
+        "--comment", type=str, default=None,
+        help="Comment for --approve/--reject"
+    )
+
+    # Output format
+    parser.add_argument(
+        "--json", action="store_true",
+        help="Output as JSON"
+    )
+
+    args = parser.parse_args()
+
+    try:
+        if args.generate:
+            if not args.control_id:
+                parser.error("--generate requires --control-id")
+            result = generate_narrative(
+                project_id=args.project_id,
+                control_id=args.control_id,
+                narrative_type=args.narrative_type,
+                use_llm=args.use_llm,
+                db_path=args.db_path,
+            )
+            if args.json:
+                print(json.dumps(result, indent=2, default=str))
+            else:
+                print(f"Narrative ID: {result['id']}")
+                print(f"Status:       {result['status']}")
+                print(f"Control:      {result['control_id']}")
+                print(f"Type:         {result['narrative_type']}")
+                print(f"LLM draft:    {'Yes' if result['llm_draft'] else 'No'}")
+
+        elif args.submit:
+            if not args.narrative_id:
+                parser.error("--submit requires --narrative-id")
+            result = submit_for_review(
+                narrative_id=args.narrative_id,
+                db_path=args.db_path,
+            )
+            if args.json:
+                print(json.dumps(result, indent=2, default=str))
+            else:
+                print(f"Narrative {result['id']} submitted for review")
+
+        elif args.approve:
+            if not args.narrative_id:
+                parser.error("--approve requires --narrative-id")
+            if not args.reviewer:
+                parser.error("--approve requires --reviewer")
+            result = approve_narrative(
+                narrative_id=args.narrative_id,
+                reviewer=args.reviewer,
+                comment=args.comment,
+                db_path=args.db_path,
+            )
+            if args.json:
+                print(json.dumps(result, indent=2, default=str))
+            else:
+                print(f"Narrative {result['id']} approved by {result['reviewer']}")
+
+        elif args.reject:
+            if not args.narrative_id:
+                parser.error("--reject requires --narrative-id")
+            if not args.reviewer:
+                parser.error("--reject requires --reviewer")
+            if not args.comment:
+                parser.error("--reject requires --comment")
+            result = reject_narrative(
+                narrative_id=args.narrative_id,
+                reviewer=args.reviewer,
+                comment=args.comment,
+                db_path=args.db_path,
+            )
+            if args.json:
+                print(json.dumps(result, indent=2, default=str))
+            else:
+                print(f"Narrative {result['id']} rejected by {result['reviewer']}")
+                print(f"Reason: {result['comment']}")
+
+        elif args.pending:
+            result = list_pending(
+                project_id=args.project_id,
+                db_path=args.db_path,
+            )
+            if args.json:
+                print(json.dumps(result, indent=2, default=str))
+            else:
+                if not result["pending"]:
+                    print("No narratives pending review.")
+                else:
+                    print(f"Pending narratives ({result['count']}):")
+                    for p in result["pending"]:
+                        print(f"  {p['id']}  {p['control_id']}  "
+                              f"({p['narrative_type']})  {p['updated_at']}")
+
+        elif args.batch:
+            result = batch_generate(
+                project_id=args.project_id,
+                use_llm=args.use_llm,
+                db_path=args.db_path,
+            )
+            if args.json:
+                print(json.dumps(result, indent=2, default=str))
+            else:
+                print(f"Generated: {result['generated']}  "
+                      f"Skipped: {result['skipped']}")
+
+    except (FileNotFoundError, ValueError) as e:
+        print(f"ERROR: {e}", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()