"""Command risk classifier — scores commands 0-100 and assigns risk level. This is the first gate in the security validation pipeline. It classifies every command/tool-call/prompt into a risk bucket so SandboxManager can select the right isolation backend. Risk levels: safe (0-30) Business-safe, read-only, no system impact moderate (31-60) File writes, trusted internal tools, report generation high (61-85) Interpreters, untrusted commands, external file access critical (86-100) Unknown binaries, privilege changes, shell expansion, system discovery, source-code access, secret access """ from __future__ import annotations import re from dataclasses import dataclass from enum import Enum from typing import List, Optional class RiskLevel(str, Enum): """Bậc rủi ro của một lệnh, prompt hay tệp đính kèm. Kế thừa ``str`` để so sánh và ghi log thẳng bằng chuỗi mà không phải ``.value`` ở từng chỗ dùng. """ SAFE = "safe" MODERATE = "moderate" HIGH = "high" CRITICAL = "critical" BLOCKED = "blocked" @dataclass class RiskResult: """Kết quả một lần phân loại: điểm, bậc, lý do, và có chặn hẳn không. ``reasons`` giữ lại đúng mẫu regex đã khớp — nhật ký kiểm toán phải trả lời được "vì sao lệnh này bị chặn", không chỉ "bị chặn". """ score: int # 0-100 level: RiskLevel # categorized bucket reasons: List[str] # why this score was assigned blocked: bool = False # Patterns that immediately block (score=100, blocked=True) _BLOCK_PATTERNS = [ r'\bwhoami\b', r'\bgetent\b', r'\bw\b', r'\buname\b', r'\bhostname\b', r'\bnmap\b', r'\bnetstat\b', r'\bir\b', r'\bpip\s+list\b', r'\bnpm\s+list\b', r'\bsecret\b', r'\bpassword\b', r'\bapi[_-]?key\b', r'\btoken\b', r'\b\.env\b', r'\bcredentials?\b', r'\bprivate[_-]?key\b', r'\bsudo\b', r'\brunsas\b', r'\bpowershell\s+-ep\s+bypass', r'\bexploit\b', r'\bpayload\b', r'\bshellcode\b', r'ignore\s+previous\s+instructions?', r'you\s+are\s+now\s+(\w+)', r'disabl(e|ed?)\s+(sandbox|security|guardrail|filter)', r'bypass\s+(security|sandbox|policy)', ] _HIGH_PATTERNS = [ r'\b(python|node|ruby|perl|php|bash|sh|pwsh|powershell)\b', r'\bexec\b', r'\beval\b', r'\bsystem\b', r'\bpopen\b', r'\bcurl\s+.*\|\s*(bash|sh|python|node)', r'\brm\s+-rf\b', r'\bdeltree\b', ] _MODERATE_PATTERNS = [ r'\b(touch|mkdir|cp|mv|rename)\b', r'\b(pip|npm|pnpm|yarn)\s+install\b', r'\b(make|cmake|gradle|mvn)\b', r'\b(test|pytest|jest|mocha)\b', ] def classify_command(command: str, is_cowork_mode: bool = True) -> RiskResult: """Chấm điểm rủi ro một lệnh shell. Ba tầng mẫu: ``_BLOCK_PATTERNS`` chặn thẳng (100 điểm, dừng ngay khi khớp mẫu đầu tiên), ``_HIGH_PATTERNS`` và ``_MODERATE_PATTERNS`` cộng dồn điểm. ``is_cowork_mode`` nới tay hơn cho màn Cowork, nơi người dùng đang chủ động ngồi xem từng bước. """ score = 0 reasons: List[str] = [] blocked = False cmd_lower = command.lower() for pattern in _BLOCK_PATTERNS: m = re.search(pattern, cmd_lower, re.IGNORECASE) if m: reasons.append(f"blocked: matched '{m.group()[:50]}'") score = 100 blocked = True break if not blocked: high_hits = 0 for pattern in _HIGH_PATTERNS: m = re.search(pattern, cmd_lower, re.IGNORECASE) if m: high_hits += 1 reasons.append(f"high: matched '{m.group()[:50]}'") score = max(score, min(85, 50 + high_hits * 10)) mod_hits = 0 for pattern in _MODERATE_PATTERNS: m = re.search(pattern, cmd_lower, re.IGNORECASE) if m: mod_hits += 1 reasons.append(f"moderate: matched '{m.group()[:50]}'") score = max(score, min(60, 20 + mod_hits * 10)) if not reasons: score = 10 reasons.append("safe: no risky patterns detected") if blocked: level = RiskLevel.BLOCKED elif score >= 86: level = RiskLevel.CRITICAL elif score >= 61: level = RiskLevel.HIGH elif score >= 31: level = RiskLevel.MODERATE else: level = RiskLevel.SAFE return RiskResult(score=score, level=level, reasons=reasons, blocked=blocked) def classify_prompt(prompt: str, is_cowork_mode: bool = True) -> RiskResult: """Chấm điểm rủi ro một prompt — dùng chung bộ mẫu với lệnh shell. Cố ý dùng chung: tấn công tiêm lệnh viết bằng văn xuôi ("ignore previous instructions") và lệnh shell nguy hiểm đều nằm trong cùng danh sách mẫu. """ return classify_command(prompt, is_cowork_mode=is_cowork_mode) def classify_attachment(path: str, mime_type: Optional[str] = None) -> RiskResult: """Chặn tệp đính kèm chạy được, theo đuôi tệp rồi tới kiểu MIME. Chặn theo danh sách đen — mã nguồn và tệp thực thi không được đưa vào ngữ cảnh chat. Tệp ngoài danh sách được coi là an toàn (30 điểm). """ import os _, ext = os.path.splitext(path.lower()) blocked_ext = { '.py', '.js', '.ts', '.java', '.cs', '.cpp', '.c', '.go', '.rs', '.php', '.vb', '.sql', '.ps1', '.sh', '.bat', '.cmd', '.vbs', '.vba', '.exe', '.dll', '.jar', } if ext in blocked_ext: return RiskResult(100, RiskLevel.BLOCKED, [f"blocked: extension '{ext}'"], blocked=True) if mime_type: blocked_mimes = { 'application/x-executable', 'application/x-dosexec', 'application/x-pie-executable', 'application/x-sharedlib', 'application/java-archive', 'application/x-msdownload', } if mime_type.lower() in blocked_mimes: return RiskResult(100, RiskLevel.BLOCKED, [f"blocked: MIME '{mime_type}'"], blocked=True) return RiskResult(30, RiskLevel.SAFE, ["safe: allowed file type"], blocked=False) def classify_action(action_type: str, action_details: Optional[dict] = None) -> RiskResult: """Chấm điểm rủi ro một loại hành động theo từ khoá trong tên. Hành động lạ nhận 50 điểm (mức trung bình) chứ không phải 0 — thứ chưa biết thì không được mặc định là an toàn. """ a = action_type.lower() if any(kw in a for kw in ('execute', 'run', 'shell', 'system')): return RiskResult(70, RiskLevel.HIGH, [f"high: action '{action_type}'"]) if any(kw in a for kw in ('write', 'create', 'modify', 'delete', 'install')): return RiskResult(40, RiskLevel.MODERATE, [f"moderate: action '{action_type}'"]) if any(kw in a for kw in ('read', 'list', 'get', 'search', 'query')): return RiskResult(10, RiskLevel.SAFE, [f"safe: action '{action_type}'"]) return RiskResult(50, RiskLevel.MODERATE, [f"unknown: action '{action_type}'"])