Governance Frameworks for AI Code Generation

from dataclasses import dataclass, field
from enum import Enum
from typing import Optional
 
class GovernanceDomain(Enum):
    TOOL_MANAGEMENT = "tool_management"
    DATA_PROTECTION = "data_protection"
    CODE_QUALITY = "code_quality"
    SECURITY = "security"
    COMPLIANCE = "compliance"
    OPERATIONS = "operations"
 
@dataclass
class GovernanceControl:
    """A specific governance control within the framework."""
    control_id: str
    domain: GovernanceDomain
    name: str
    description: str
    implementation: str
    evidence: list[str]
    frequency: str
 
GOVERNANCE_CONTROLS = [
    # Tool Management
    GovernanceControl(
        control_id="TM-001",
        domain=GovernanceDomain.TOOL_MANAGEMENT,
        name="AI Tool Inventory",
        description="Maintain a comprehensive inventory of all AI coding tools in use",
        implementation=(
            "Deploy endpoint detection to identify AI coding tools. "
            "Maintain a registry of approved tools with version tracking."
        ),
        evidence=[
            "Tool inventory spreadsheet updated monthly",
            "EDR detection rules for AI coding tools",
            "Software asset management integration",
        ],
        frequency="Monthly",
    ),
    GovernanceControl(
        control_id="TM-002",
        domain=GovernanceDomain.TOOL_MANAGEMENT,
        name="AI Tool Approval Process",
        description="Formal approval process for new AI coding tools",
        implementation=(
            "Security and legal review required before any new AI coding "
            "tool is approved for use. Review covers: data handling, "
            "terms of service, security architecture, and compliance."
        ),
        evidence=[
            "Approval request forms with security/legal sign-off",
            "Tool security assessment reports",
            "Legal review of terms of service",
        ],
        frequency="Per new tool request",
    ),
    GovernanceControl(
        control_id="TM-003",
        domain=GovernanceDomain.TOOL_MANAGEMENT,
        name="AI Tool Configuration Standards",
        description="Standard security configurations for approved AI tools",
        implementation=(
            "Define and enforce standard configurations including: "
            "telemetry settings, .cursorignore/.aiderignore templates, "
            "proxy requirements, and model endpoint restrictions."
        ),
        evidence=[
            "Configuration templates per tool",
            "Compliance scan results",
            "Developer onboarding documentation",
        ],
        frequency="Quarterly review",
    ),
 
    # Data Protection
    GovernanceControl(
        control_id="DP-001",
        domain=GovernanceDomain.DATA_PROTECTION,
        name="Data Classification for AI Tools",
        description="Classify code repositories by sensitivity for AI tool access",
        implementation=(
            "Apply data classification labels to all repositories. "
            "Restrict AI tool usage based on classification: "
            "PUBLIC=all tools, INTERNAL=approved tools only, "
            "CONFIDENTIAL=enterprise endpoints only, "
            "SECRET=no AI tools."
        ),
        evidence=[
            "Repository classification inventory",
            "AI tool access matrix by classification",
            "Classification enforcement scan results",
        ],
        frequency="Quarterly",
    ),
    GovernanceControl(
        control_id="DP-002",
        domain=GovernanceDomain.DATA_PROTECTION,
        name="Code Transmission Monitoring",
        description="Monitor code transmitted to AI service providers",
        implementation=(
            "Deploy network monitoring to track data egress to AI "
            "service endpoints. Alert on transmission of code from "
            "CONFIDENTIAL or higher repositories."
        ),
        evidence=[
            "Network monitoring configuration",
            "Alert rules for AI service endpoints",
            "Monthly data egress reports",
        ],
        frequency="Continuous",
    ),
 
    # Security
    GovernanceControl(
        control_id="SEC-001",
        domain=GovernanceDomain.SECURITY,
        name="AI Code Security Scanning",
        description="Automated security scanning of AI-generated code",
        implementation=(
            "Integrate Semgrep with AI-specific rule sets into CI/CD "
            "pipeline. All pull requests are scanned with both standard "
            "and AI-specific security rules."
        ),
        evidence=[
            "CI/CD pipeline configuration showing Semgrep integration",
            "AI-specific Semgrep rule set",
            "Monthly scan results summary",
        ],
        frequency="Every pull request",
    ),
    GovernanceControl(
        control_id="SEC-002",
        domain=GovernanceDomain.SECURITY,
        name="Prompt Injection Prevention",
        description="Controls to prevent prompt injection via repository content",
        implementation=(
            "Scan .cursorrules, CLAUDE.md, and similar configuration files "
            "for potentially dangerous directives. Require security review "
            "for changes to AI configuration files."
        ),
        evidence=[
            "Pre-commit hooks for AI config file changes",
            "CODEOWNERS rules for AI config files",
            "Scan results for prompt injection patterns",
        ],
        frequency="Every commit (automated)",
    ),
    GovernanceControl(
        control_id="SEC-003",
        domain=GovernanceDomain.SECURITY,
        name="AI Code Review Requirements",
        description="Enhanced code review requirements for AI-generated code",
        implementation=(
            "Pull requests identified as containing AI-generated code "
            "require additional review focusing on: input validation, "
            "authorization, error handling, and secure defaults."
        ),
        evidence=[
            "Code review checklist for AI-generated code",
            "PR review metrics showing compliance",
            "Reviewer training completion records",
        ],
        frequency="Every pull request",
    ),
 
    # Compliance
    GovernanceControl(
        control_id="CMP-001",
        domain=GovernanceDomain.COMPLIANCE,
        name="License Compliance Scanning",
        description="Scan AI-generated code for license contamination",
        implementation=(
            "Deploy ScanCode Toolkit or similar tool to scan for "
            "copyleft license indicators in AI-generated code. "
            "Block merges with unresolved license findings."
        ),
        evidence=[
            "License scan integration in CI/CD",
            "Monthly compliance reports",
            "Remediation records for flagged code",
        ],
        frequency="Every pull request",
    ),
    GovernanceControl(
        control_id="CMP-002",
        domain=GovernanceDomain.COMPLIANCE,
        name="Regulatory Compliance Assessment",
        description="Assess AI code generation against regulatory requirements",
        implementation=(
            "Map AI code generation practices against applicable regulations "
            "(SOC 2, HIPAA, PCI DSS, GDPR). Document how AI tool usage "
            "affects compliance posture."
        ),
        evidence=[
            "Regulatory mapping document",
            "Compliance gap analysis",
            "Auditor-ready documentation",
        ],
        frequency="Annual + per regulation change",
    ),
 
    # Operations
    GovernanceControl(
        control_id="OPS-001",
        domain=GovernanceDomain.OPERATIONS,
        name="Developer Training Program",
        description="Training program for secure use of AI coding tools",
        implementation=(
            "Mandatory training for all developers using AI coding tools. "
            "Covers: secure prompting, code review for AI output, "
            "data handling obligations, and incident reporting."
        ),
        evidence=[
            "Training curriculum and materials",
            "Completion records per developer",
            "Assessment scores",
        ],
        frequency="Annual + on tool onboarding",
    ),
    GovernanceControl(
        control_id="OPS-002",
        domain=GovernanceDomain.OPERATIONS,
        name="Incident Response for AI Code Issues",
        description="Incident response procedures for AI code generation incidents",
        implementation=(
            "Documented playbook for responding to: data exfiltration "
            "via AI tools, license violations in AI-generated code, "
            "vulnerabilities introduced by AI, and compromised AI tools."
        ),
        evidence=[
            "Incident response playbook",
            "Tabletop exercise records",
            "Post-incident review reports",
        ],
        frequency="Annual playbook review + per incident",
    ),
]

@dataclass
class RiskEntry:
    """Risk register entry for AI code generation."""
    risk_id: str
    category: str
    description: str
    likelihood: str      # rare, unlikely, possible, likely, almost_certain
    impact: str          # negligible, minor, moderate, major, severe
    inherent_risk: str   # low, medium, high, critical
    controls: list[str]  # Control IDs that mitigate this risk
    residual_risk: str   # Risk level after controls
 
RISK_REGISTER = [
    RiskEntry(
        risk_id="R001",
        category="Data Protection",
        description="Proprietary source code transmitted to AI service provider and exposed through data breach or training",
        likelihood="possible",
        impact="major",
        inherent_risk="high",
        controls=["DP-001", "DP-002", "TM-002"],
        residual_risk="medium",
    ),
    RiskEntry(
        risk_id="R002",
        category="Security",
        description="AI-generated code introduces critical vulnerability (SQLi, RCE) into production",
        likelihood="likely",
        impact="major",
        inherent_risk="critical",
        controls=["SEC-001", "SEC-003"],
        residual_risk="medium",
    ),
    RiskEntry(
        risk_id="R003",
        category="Compliance",
        description="AI-generated code contains GPL-licensed material in proprietary product",
        likelihood="possible",
        impact="major",
        inherent_risk="high",
        controls=["CMP-001"],
        residual_risk="medium",
    ),
    RiskEntry(
        risk_id="R004",
        category="Security",
        description="Prompt injection via repository content leads to malicious code execution",
        likelihood="unlikely",
        impact="severe",
        inherent_risk="high",
        controls=["SEC-002", "SEC-003"],
        residual_risk="medium",
    ),
    RiskEntry(
        risk_id="R005",
        category="Supply Chain",
        description="AI suggests hallucinated or typosquatted package that is malicious",
        likelihood="possible",
        impact="major",
        inherent_risk="high",
        controls=["SEC-001", "TM-003"],
        residual_risk="medium",
    ),
    RiskEntry(
        risk_id="R006",
        category="Operational",
        description="Developers over-rely on AI tools and miss security issues they would have caught manually",
        likelihood="likely",
        impact="moderate",
        inherent_risk="high",
        controls=["OPS-001", "SEC-003"],
        residual_risk="medium",
    ),
    RiskEntry(
        risk_id="R007",
        category="Data Protection",
        description="AI tool accesses and transmits secrets (API keys, credentials) from developer environment",
        likelihood="possible",
        impact="severe",
        inherent_risk="critical",
        controls=["TM-003", "DP-001"],
        residual_risk="medium",
    ),
    RiskEntry(
        risk_id="R008",
        category="Compliance",
        description="AI code generation practices violate regulatory requirements (SOC 2, HIPAA)",
        likelihood="possible",
        impact="major",
        inherent_risk="high",
        controls=["CMP-002", "TM-002"],
        residual_risk="low",
    ),
]
 
def generate_risk_heatmap(risks: list[RiskEntry]) -> str:
    """Generate a text-based risk heatmap."""
    heatmap = "Risk Heatmap (Inherent Risk)\n"
    heatmap += "=" * 60 + "\n\n"
 
    likelihood_order = ["almost_certain", "likely", "possible", "unlikely", "rare"]
    impact_order = ["negligible", "minor", "moderate", "major", "severe"]
 
    heatmap += f"{'':>15} | {'Negligible':>10} | {'Minor':>10} | {'Moderate':>10} | {'Major':>10} | {'Severe':>10}\n"
    heatmap += "-" * 80 + "\n"
 
    for likelihood in likelihood_order:
        row_risks = [r.risk_id for r in risks if r.likelihood == likelihood]
        cells = []
        for impact in impact_order:
            matching = [
                r.risk_id for r in risks
                if r.likelihood == likelihood and r.impact == impact
            ]
            cells.append(", ".join(matching) if matching else "-")
 
        heatmap += f"{likelihood:>15} | {cells[0]:>10} | {cells[1]:>10} | {cells[2]:>10} | {cells[3]:>10} | {cells[4]:>10}\n"
 
    return heatmap

@dataclass
class MaturityLevel:
    level: int
    name: str
    description: str
    characteristics: list[str]
    controls_required: list[str]
 
MATURITY_LEVELS = [
    MaturityLevel(
        level=1,
        name="Initial",
        description="Ad hoc AI tool usage with no governance",
        characteristics=[
            "Developers use AI tools without approval",
            "No visibility into which tools are used or what code they generate",
            "No specific security controls for AI-generated code",
            "No data classification for AI tool access",
            "No training on AI coding tool risks",
        ],
        controls_required=[],
    ),
    MaturityLevel(
        level=2,
        name="Developing",
        description="Basic awareness and initial controls",
        characteristics=[
            "Approved AI tool list exists",
            "Developers aware of AI coding risks",
            "Basic .cursorignore/.aiderignore templates distributed",
            "Standard code review process applies to AI code",
            "Incident response plan acknowledges AI code risks",
        ],
        controls_required=["TM-001", "TM-002", "OPS-001"],
    ),
    MaturityLevel(
        level=3,
        name="Defined",
        description="Formal policies and automated controls",
        characteristics=[
            "Formal AI code generation policy documented",
            "Automated security scanning with AI-specific rules",
            "Data classification enforced for AI tool access",
            "License compliance scanning in CI/CD",
            "AI configuration files (CLAUDE.md, .cursorrules) reviewed",
            "Developer training program with assessment",
        ],
        controls_required=[
            "TM-001", "TM-002", "TM-003",
            "DP-001", "DP-002",
            "SEC-001", "SEC-002", "SEC-003",
            "CMP-001",
            "OPS-001", "OPS-002",
        ],
    ),
    MaturityLevel(
        level=4,
        name="Managed",
        description="Metrics-driven governance with continuous improvement",
        characteristics=[
            "AI code generation metrics tracked (volume, vulnerability rate, review coverage)",
            "Risk register maintained and reviewed quarterly",
            "Regular tabletop exercises for AI code incidents",
            "Automated enforcement of all policies",
            "Supply chain monitoring for AI-suggested dependencies",
            "Regulatory compliance assessed and documented",
        ],
        controls_required=[
            "TM-001", "TM-002", "TM-003",
            "DP-001", "DP-002",
            "SEC-001", "SEC-002", "SEC-003",
            "CMP-001", "CMP-002",
            "OPS-001", "OPS-002",
        ],
    ),
    MaturityLevel(
        level=5,
        name="Optimizing",
        description="Advanced governance with proactive risk management",
        characteristics=[
            "Predictive risk analysis for AI code generation trends",
            "Automated policy generation based on threat intelligence",
            "AI-assisted governance (using AI to monitor AI code quality)",
            "Industry benchmarking and best practice leadership",
            "Continuous red team exercises targeting AI code generation",
            "Zero-trust architecture for AI tool execution",
        ],
        controls_required=[
            "All controls plus advanced monitoring and analytics",
        ],
    ),
]
 
def assess_maturity(implemented_controls: list[str]) -> dict:
    """Assess organizational maturity based on implemented controls."""
    current_level = 1
    for level in MATURITY_LEVELS:
        if all(
            ctrl in implemented_controls
            for ctrl in level.controls_required
        ):
            current_level = level.level
        else:
            break
 
    current = next(l for l in MATURITY_LEVELS if l.level == current_level)
    next_level = next(
        (l for l in MATURITY_LEVELS if l.level == current_level + 1), None
    )
 
    gaps = []
    if next_level:
        gaps = [
            ctrl for ctrl in next_level.controls_required
            if ctrl not in implemented_controls
        ]
 
    return {
        "current_level": current_level,
        "level_name": current.name,
        "description": current.description,
        "next_level": next_level.level if next_level else None,
        "gaps_to_next_level": gaps,
    }

def generate_usage_policy(org_name: str, approved_tools: list[str]) -> str:
    """Generate an AI code generation usage policy."""
    tools_list = "\n".join(f"  - {tool}" for tool in approved_tools)
 
    return f"""
# {org_name} — AI Code Generation Usage Policy
 
## Purpose
This policy governs the use of AI code generation tools within {org_name}
to manage security, compliance, and operational risk.
 
## Scope
This policy applies to all employees, contractors, and third parties who
write code for {org_name} systems.
 
## Approved Tools
The following AI code generation tools are approved for use:
{tools_list}
 
Use of unapproved tools is prohibited. To request approval for a new tool,
submit a request to the Security Engineering team.
 
## Data Classification Restrictions
- PUBLIC repositories: All approved tools permitted
- INTERNAL repositories: Approved tools with enterprise endpoints only
- CONFIDENTIAL repositories: Approved tools with enterprise endpoints,
  privacy mode enabled, no telemetry
- SECRET/RESTRICTED repositories: AI code generation tools are PROHIBITED
 
## Required Controls
1. All repositories using AI tools MUST have appropriate ignore files
   (.cursorignore, .aiderignore) to exclude sensitive content
2. All AI-generated code MUST pass automated security scanning before merge
3. Pull requests with AI-generated code MUST follow the enhanced review
   checklist
4. API keys for AI tools MUST be provisioned through corporate key
   management, not personal accounts
5. Agent/autonomous execution modes MUST be restricted to sandboxed
   environments
 
## Developer Responsibilities
- Complete AI code generation security training before using any AI tool
- Review all AI-generated code for security issues before accepting
- Report suspected security incidents involving AI tools immediately
- Do not share corporate API keys or use personal API keys for work
- Do not disable security controls or telemetry restrictions
 
## Monitoring and Enforcement
- AI tool usage is monitored through endpoint detection and network analysis
- Violations of this policy may result in revocation of AI tool access
- Repeat violations will be escalated per the Employee Conduct Policy
 
## Review
This policy is reviewed quarterly by the Security Engineering team.
Last reviewed: {datetime.utcnow().strftime('%Y-%m-%d')}
"""

AI_GOVERNANCE_KPIS = {
    "adoption_metrics": [
        {
            "kpi": "AI Tool Coverage",
            "definition": "Percentage of developers with approved AI tools configured",
            "target": ">90%",
            "measurement": "EDR inventory / developer headcount",
        },
        {
            "kpi": "Configuration Compliance",
            "definition": "Percentage of AI tool installations meeting configuration standards",
            "target": ">95%",
            "measurement": "Compliance scan passing / total installations",
        },
    ],
    "security_metrics": [
        {
            "kpi": "AI Vulnerability Detection Rate",
            "definition": "Percentage of AI-introduced vulnerabilities caught before production",
            "target": ">99% for critical, >95% for high",
            "measurement": "Pre-production findings / (pre-production + post-production findings)",
        },
        {
            "kpi": "Mean Time to Detect AI Vulnerability",
            "definition": "Average time from code generation to vulnerability detection",
            "target": "<24 hours",
            "measurement": "Detection timestamp - commit timestamp",
        },
        {
            "kpi": "AI Security Scan Coverage",
            "definition": "Percentage of PRs with AI code scanned by AI-specific rules",
            "target": "100%",
            "measurement": "PRs with AI scan / total PRs with AI-generated code",
        },
    ],
    "compliance_metrics": [
        {
            "kpi": "License Compliance Rate",
            "definition": "Percentage of AI-generated code passing license compliance scan",
            "target": "100%",
            "measurement": "Passing scans / total scans",
        },
        {
            "kpi": "Data Classification Compliance",
            "definition": "Percentage of repositories with correct AI tool access controls",
            "target": ">98%",
            "measurement": "Compliant repos / total repos",
        },
    ],
    "training_metrics": [
        {
            "kpi": "Training Completion Rate",
            "definition": "Percentage of AI tool users who completed security training",
            "target": "100%",
            "measurement": "Trained users / total AI tool users",
        },
        {
            "kpi": "Training Assessment Score",
            "definition": "Average score on AI code security assessment",
            "target": ">80%",
            "measurement": "Mean assessment score",
        },
    ],
}

IMPLEMENTATION_ROADMAP = {
    "month_1_2": {
        "name": "Foundation",
        "activities": [
            "Establish AI code generation governance committee",
            "Conduct tool inventory across the organization",
            "Draft initial AI code generation usage policy",
            "Deploy AI-specific Semgrep rules in CI/CD",
            "Create .cursorignore and CLAUDE.md templates",
        ],
        "target_maturity": 2,
    },
    "month_3_4": {
        "name": "Formalization",
        "activities": [
            "Implement data classification for AI tool access",
            "Deploy network monitoring for AI service endpoints",
            "Launch developer training program",
            "Implement license compliance scanning",
            "Establish incident response playbook",
        ],
        "target_maturity": 3,
    },
    "month_5_6": {
        "name": "Measurement",
        "activities": [
            "Deploy governance metrics dashboard",
            "Conduct first tabletop exercise",
            "Complete regulatory compliance assessment",
            "Review and refine all policies based on metrics",
            "Assess and address gaps for Level 4",
        ],
        "target_maturity": 3,  # Solidify Level 3
    },
    "month_7_12": {
        "name": "Optimization",
        "activities": [
            "Implement predictive risk analytics",
            "Automate policy enforcement",
            "Conduct red team exercise targeting AI code generation",
            "Benchmark against industry peers",
            "Plan for Level 5 capabilities",
        ],
        "target_maturity": 4,
    },
}