"""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())