app/tools/file_scanner.py

"""
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:
        <source_file name="path/to/file.py">
        ...content...
        </source_file>
    """
    safe_name = filename.replace('"', "'")  # prevent attribute injection
    return f'<source_file name="{safe_name}">\n{content}\n</source_file>'


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]