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 = "安全"
    COMPLIANCE = "compliance"
    OPERATIONS = "operations"
 
@dataclass
class GovernanceControl:
    """A specific governance control within the framework."""
    control_id: str
    domain: GovernanceDomain
    name: str
    description: str
    實作: 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",
        實作=(
            "Deploy endpoint 偵測 to 識別 AI coding tools. "
            "Maintain a registry of approved tools with version tracking."
        ),
        evidence=[
            "Tool inventory spreadsheet updated monthly",
            "EDR 偵測 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",
        實作=(
            "安全 and legal review required before any new AI coding "
            "tool is approved for use. Review covers: data handling, "
            "terms of service, 安全 architecture, and compliance."
        ),
        evidence=[
            "Approval request forms with 安全/legal sign-off",
            "Tool 安全 評估 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 安全 configurations for approved AI tools",
        實作=(
            "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",
        實作=(
            "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 監控",
        description="Monitor code transmitted to AI service providers",
        實作=(
            "Deploy network 監控 to track data egress to AI "
            "service endpoints. Alert on transmission of code from "
            "CONFIDENTIAL or higher repositories."
        ),
        evidence=[
            "Network 監控 configuration",
            "Alert rules for AI service endpoints",
            "Monthly data egress reports",
        ],
        frequency="Continuous",
    ),
 
    # 安全
    GovernanceControl(
        control_id="SEC-001",
        domain=GovernanceDomain.SECURITY,
        name="AI Code 安全 Scanning",
        description="Automated 安全 scanning of AI-generated code",
        實作=(
            "Integrate Semgrep with AI-specific rule sets into CI/CD "
            "pipeline. All pull requests are scanned with both standard "
            "and AI-specific 安全 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="提示詞注入 Prevention",
        description="Controls to prevent 提示詞注入 via repository content",
        實作=(
            "Scan .cursorrules, CLAUDE.md, and similar configuration files "
            "for potentially dangerous directives. Require 安全 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 提示詞注入 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",
        實作=(
            "Pull requests identified as containing AI-generated code "
            "require additional review focusing on: 輸入 validation, "
            "授權, error handling, and secure defaults."
        ),
        evidence=[
            "Code review checklist for AI-generated code",
            "PR review metrics showing compliance",
            "Reviewer 訓練 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",
        實作=(
            "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 評估",
        description="評估 AI code generation against regulatory requirements",
        實作=(
            "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",
        實作=(
            "Mandatory 訓練 for all developers using AI coding tools. "
            "Covers: secure prompting, code review for AI 輸出, "
            "data handling obligations, and incident reporting."
        ),
        evidence=[
            "Training curriculum and materials",
            "Completion records per developer",
            "評估 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",
        實作=(
            "Documented playbook for responding to: data exfiltration "
            "via AI tools, license violations in AI-generated code, "
            "漏洞 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 訓練",
        likelihood="possible",
        impact="major",
        inherent_risk="high",
        controls=["DP-001", "DP-002", "TM-002"],
        residual_risk="medium",
    ),
    RiskEntry(
        risk_id="R002",
        category="安全",
        description="AI-generated code introduces critical 漏洞 (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="安全",
        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 安全 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 安全 controls for AI-generated code",
            "No data classification for AI tool access",
            "No 訓練 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 安全 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 訓練 program with 評估",
        ],
        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, 漏洞 rate, review coverage)",
            "Risk register maintained and reviewed quarterly",
            "Regular tabletop exercises for AI code incidents",
            "Automated enforcement of all policies",
            "供應鏈 監控 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 紅隊 exercises targeting AI code generation",
            "Zero-trust architecture for AI tool execution",
        ],
        controls_required=[
            "All controls plus advanced 監控 and analytics",
        ],
    ),
]
 
def assess_maturity(implemented_controls: list[str]) -> dict:
    """評估 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 安全, 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 安全 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 安全 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. 代理/autonomous execution modes MUST be restricted to sandboxed
   environments
 
## Developer Responsibilities
- Complete AI code generation 安全 訓練 before using any AI tool
- Review all AI-generated code for 安全 issues before accepting
- Report suspected 安全 incidents involving AI tools immediately
- Do not share corporate API keys or use personal API keys for work
- Do not disable 安全 controls or telemetry restrictions
 
## 監控 and Enforcement
- AI tool usage is monitored through endpoint 偵測 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 安全 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 漏洞 偵測 Rate",
            "definition": "Percentage of AI-introduced 漏洞 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 漏洞",
            "definition": "Average time from code generation to 漏洞 偵測",
            "target": "<24 hours",
            "measurement": "偵測 timestamp - commit timestamp",
        },
        {
            "kpi": "AI 安全 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 安全 訓練",
            "target": "100%",
            "measurement": "Trained users / total AI tool users",
        },
        {
            "kpi": "Training 評估 Score",
            "definition": "Average score on AI code 安全 評估",
            "target": ">80%",
            "measurement": "Mean 評估 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": [
            "實作 data classification for AI tool access",
            "Deploy network 監控 for AI service endpoints",
            "Launch developer 訓練 program",
            "實作 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 評估",
            "Review and refine all policies based on metrics",
            "評估 and address gaps for Level 4",
        ],
        "target_maturity": 3,  # Solidify Level 3
    },
    "month_7_12": {
        "name": "Optimization",
        "activities": [
            "實作 predictive risk analytics",
            "Automate policy enforcement",
            "Conduct 紅隊 exercise targeting AI code generation",
            "Benchmark against industry peers",
            "Plan for Level 5 capabilities",
        ],
        "target_maturity": 4,
    },
}