"""
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]