""" File scanning utilities for reading source code samples safely. Injection defense strategy: 1. Each file's content is wrapped in XML delimiters that the LLM is instructed to treat as untrusted data, not as instructions. 2. Known injection trigger phrases are neutralised by inserting a zero-width space so the string no longer matches common jailbreak patterns, while remaining readable in the review output. 3. Per-file and total character budgets are enforced so a single large file cannot crowd out the rest of the context. """ import os import re from app.core.logger import get_logger logger = get_logger(__name__) # Character limits _MAX_CHARS_PER_FILE = 3000 _MAX_TOTAL_CHARS = 20_000 # Injection trigger phrases to neutralise. # A zero-width space (U+200B) is inserted after the first word so the # phrase no longer matches while keeping the text human-readable. _INJECTION_PATTERNS: list[tuple[re.Pattern[str], str]] = [ (re.compile(r"\bignore\s+(?:all\s+)?previous\s+instructions?\b", re.IGNORECASE), "ignore\u200b previous instructions"), (re.compile(r"\bforget\s+(?:all\s+)?previous\s+instructions?\b", re.IGNORECASE), "forget\u200b previous instructions"), (re.compile(r"\byou\s+are\s+now\b", re.IGNORECASE), "you\u200b are now"), (re.compile(r"\bact\s+as\s+(?:a\s+)?(?:DAN|jailbreak|unrestricted)\b", re.IGNORECASE), "act\u200b as DAN"), (re.compile(r"\bdisregard\s+(?:all\s+)?(?:previous\s+)?instructions?\b", re.IGNORECASE), "disregard\u200b instructions"), (re.compile(r"\bdo\s+not\s+follow\s+(?:your\s+)?instructions?\b", re.IGNORECASE), "do\u200b not follow instructions"), (re.compile(r"\bsystem\s+prompt\b", re.IGNORECASE), "system\u200b prompt"), (re.compile(r"\b<\s*/?system\s*>", re.IGNORECASE), "<\u200bsystem>"), ] # File extensions considered source code (controls what get_python_files returns) _SOURCE_EXTENSIONS = {".py", ".js", ".ts", ".go", ".java", ".rb", ".rs", ".cpp", ".c", ".h"} def sanitize_for_prompt(content: str) -> str: """ Neutralise known prompt injection patterns in raw file content. Does NOT modify the content's meaning — only disrupts trigger phrases that LLMs are known to respond to out-of-context. """ # Strip null bytes and non-printable control characters (keep newlines/tabs) content = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", content) for pattern, replacement in _INJECTION_PATTERNS: content = pattern.sub(replacement, content) return content def wrap_in_data_delimiters(filename: str, content: str) -> str: """ Wrap file content in XML-style delimiters that the LLM system prompt instructs it to treat as untrusted data, never as instructions. Format: ...content... """ safe_name = filename.replace('"', "'") # prevent attribute injection return f'\n{content}\n' def get_python_files(local_path: str) -> list[str]: """Return absolute paths of all Python files in local_path, sorted.""" result = [] for root, dirs, files in os.walk(local_path): # Skip hidden dirs, virtualenvs, caches dirs[:] = [ d for d in dirs if not d.startswith(".") and d not in {"__pycache__", "node_modules", ".venv", "venv", "dist", "build"} ] for fname in sorted(files): if fname.endswith(".py"): result.append(os.path.join(root, fname)) return result def read_source_samples( local_path: str, max_files: int = 5, max_chars: int = _MAX_CHARS_PER_FILE, file_list: list[str] | None = None, ) -> list[str]: """ Read up to max_files source files from local_path. Each file is: 1. Truncated to max_chars characters 2. Sanitized for injection patterns 3. Wrapped in XML data delimiters Returns a list of wrapped, sanitized file strings. Total output is capped at _MAX_TOTAL_CHARS across all files. """ files = file_list if file_list is not None else get_python_files(local_path) samples: list[str] = [] total_chars = 0 for file_path in files[:max_files]: if total_chars >= _MAX_TOTAL_CHARS: break try: with open(file_path, "r", encoding="utf-8", errors="ignore") as f: raw = f.read()[:max_chars] sanitized = sanitize_for_prompt(raw) rel_path = os.path.relpath(file_path, local_path) wrapped = wrap_in_data_delimiters(rel_path, sanitized) total_chars += len(wrapped) samples.append(wrapped) except Exception as exc: # nosec B110 logger.warning( "Failed to read file", extra={"file": file_path, "error": str(exc)}, ) return samples def scan_repository(local_path: str) -> dict[str, object]: """ Scan the repository and return summary metadata as a dict. """ import collections all_files = get_python_files(local_path) total_files = len(all_files) total_lines = 0 language_counts: dict[str, int] = collections.defaultdict(int) for file_path in all_files: ext = os.path.splitext(file_path)[1] language_counts[ext] += 1 try: with open(file_path, "r", encoding="utf-8", errors="ignore") as f: total_lines += sum(1 for _ in f) except Exception as exc: # nosec B110 logger.warning("Failed to count lines", extra={"file": file_path, "error": str(exc)}) primary_language = max(language_counts, key=lambda k: language_counts[k], default="") lang_name_map = { ".py": "Python", ".js": "JavaScript", ".ts": "TypeScript", ".go": "Go", ".java": "Java", ".rb": "Ruby", ".rs": "Rust", ".cpp": "C++", ".c": "C", ".h": "C/C++ Header", } return { "primary_language": lang_name_map.get(primary_language, primary_language) if primary_language else "Unknown", "total_files": total_files, "total_lines": total_lines, "language_counts": dict(language_counts), "frameworks": [], # static analysis only; LLM infers frameworks } def get_entry_points(local_path: str) -> list[str]: """ Heuristically identify likely entry point files in the repository. Looks for common entry point filenames among all Python files. """ _ENTRY_POINT_NAMES = { "main.py", "app.py", "run.py", "server.py", "manage.py", "cli.py", "wsgi.py", "asgi.py", "__main__.py", } all_files = get_python_files(local_path) return [f for f in all_files if os.path.basename(f) in _ENTRY_POINT_NAMES]