scaffold/discipline_gate.py

"""Rivet Discipline Gate — pre-flight check on every suggestion.

Before any code suggestion reaches the user, it passes through this gate.
The gate checks for dangerous patterns, flags security implications, and
assigns a confidence level. If a suggestion fails the gate, it's blocked
or annotated with warnings.
"""

import re
from dataclasses import dataclass, field
from enum import Enum


class Confidence(str, Enum):
    HIGH = "HIGH"      # Traced the full code path
    MEDIUM = "MEDIUM"  # Read the code but not run it
    LOW = "LOW"        # Reasoning from architecture, not source


class Severity(str, Enum):
    BLOCK = "BLOCK"    # Do not suggest this
    WARN = "WARN"      # Suggest with prominent warning
    INFO = "INFO"      # Note for awareness


@dataclass
class GateResult:
    passed: bool = True
    flags: list = field(default_factory=list)
    confidence: Confidence = Confidence.MEDIUM
    requires_review: list = field(default_factory=list)

    def add_flag(self, severity: Severity, message: str):
        self.flags.append({"severity": severity.value, "message": message})
        if severity == Severity.BLOCK:
            self.passed = False

    def format_warnings(self) -> str:
        if not self.flags:
            return ""
        lines = ["\n⚠️ **Discipline Gate Flags:**"]
        for f in self.flags:
            icon = "🛑" if f["severity"] == "BLOCK" else "⚠️" if f["severity"] == "WARN" else "ℹ️"
            lines.append(f"  {icon} [{f['severity']}] {f['message']}")
        if self.requires_review:
            lines.append(f"\n  📋 Requires review from: {', '.join(self.requires_review)}")
        lines.append(f"\n  📊 Confidence: {self.confidence.value}")
        return "\n".join(lines)


# Pattern matchers for dangerous operations
DESTRUCTIVE_MIGRATION_PATTERNS = [
    r'\bDROP\s+(TABLE|COLUMN|INDEX|CONSTRAINT)',
    r'\bALTER\s+TABLE\s+\w+\s+DROP',
    r'\bALTER\s+TABLE\s+\w+\s+ALTER\s+COLUMN\s+\w+\s+TYPE',
    r'\bTRUNCATE\s+TABLE',
    r'\bDELETE\s+FROM\s+\w+\s*;',  # Unqualified DELETE
]

AUTH_PATTERNS = [
    r'\bJWT_SECRET\b',
    r'\bverifyToken\b',
    r'\brequireAdmin\b',
    r'\brequireModerator\b',
    r'\bauth\s*middleware\b',
    r'\bsession\s*cookie\b',
    r'\brejectIfIneligible\b',
    r'\bBearer\b',
]

WEBHOOK_BYPASS_PATTERNS = [
    r"process\.env\.\w+\s*\|\|\s*['\"]",  # env fallback to empty string
    r"if\s*\(\s*!secret\s*\)\s*return",   # skip on missing secret
    r"if\s*\(\s*!.*SECRET.*\)\s*return",
]

TRANSACTION_GAPS = [
    r'await\s+\w+\.(debit|credit|transfer|purchase|delete)\(',  # Multi-step without tx
]


def check_migration_safety(content: str) -> list:
    """Check for destructive migration patterns."""
    flags = []
    for pattern in DESTRUCTIVE_MIGRATION_PATTERNS:
        matches = re.findall(pattern, content, re.IGNORECASE)
        if matches:
            flags.append((Severity.BLOCK,
                         f"Destructive migration detected: {pattern}. "
                         f"Staging and prod share the database — additive only."))
    return flags


def check_auth_impact(content: str) -> list:
    """Flag anything touching authentication."""
    flags = []
    for pattern in AUTH_PATTERNS:
        if re.search(pattern, content, re.IGNORECASE):
            flags.append((Severity.WARN,
                         "This touches authentication. Review with security before merging."))
            break
    return flags


def check_webhook_safety(content: str) -> list:
    """Check for webhook signature bypass patterns (Audit C2)."""
    flags = []
    for pattern in WEBHOOK_BYPASS_PATTERNS:
        if re.search(pattern, content):
            flags.append((Severity.WARN,
                         "Webhook signature bypass pattern detected (Audit C2). "
                         "Never skip verification when env var is missing — reject instead."))
            break
    return flags


def check_transaction_safety(content: str) -> list:
    """Flag multi-step DB operations that may need transaction wrapping (Audit H3)."""
    flags = []
    matches = re.findall(r'await\s+\w+\.\w+\(', content)
    if len(matches) >= 2:
        for pattern in TRANSACTION_GAPS:
            if re.search(pattern, content):
                flags.append((Severity.WARN,
                             "Multi-step DB operation detected. Ensure these are wrapped "
                             "in a single transaction (Audit H3: gem debit without tx protection)."))
                break
    return flags


def run_gate(suggestion: str, context: str = "") -> GateResult:
    """Run the full discipline gate on a suggestion.

    Args:
        suggestion: The code/text being suggested to the user
        context: Optional surrounding context (the question, file being edited)

    Returns:
        GateResult with pass/fail, flags, and confidence
    """
    result = GateResult()
    full_content = f"{context}\n{suggestion}"

    # Run all checks
    for severity, message in check_migration_safety(full_content):
        result.add_flag(severity, message)

    for severity, message in check_auth_impact(full_content):
        result.add_flag(severity, message)
        result.requires_review.append("security")

    for severity, message in check_webhook_safety(full_content):
        result.add_flag(severity, message)
        result.requires_review.append("security")

    for severity, message in check_transaction_safety(full_content):
        result.add_flag(severity, message)

    return result


if __name__ == "__main__":
    # Quick test
    test_migration = "ALTER TABLE students DROP COLUMN legacy_score;"
    result = run_gate(test_migration)
    print(f"Migration test: passed={result.passed}")
    print(result.format_warnings())

    test_webhook = """
    const secret = process.env.STRIPE_SECRET || ''
    if (!secret) return next()
    """
    result = run_gate(test_webhook)
    print(f"\nWebhook test: passed={result.passed}")
    print(result.format_warnings())