ingest.py

"""Turn whatever the user gives us into a clean list of analysable files.

Handles a ZIP upload, a single pasted/uploaded file, a folder path, or a
GitHub URL. Filters out noise (node_modules, .git, binaries, lockfiles),
enforces size caps, and — critically — runs a **secret scan** so we never
display or send API keys and passwords to the model.
"""
from __future__ import annotations

import os
import re
import tempfile
import zipfile
from dataclasses import dataclass, field

import config


@dataclass
class SourceFile:
    path: str          # project-relative, forward slashes
    abspath: str       # where it lives on disk right now
    text: str          # file contents, with secrets already redacted
    lang: str


@dataclass
class Ingested:
    name: str
    files: list[SourceFile] = field(default_factory=list)
    skipped: list[str] = field(default_factory=list)       # listed but not parsed
    secrets_found: list[str] = field(default_factory=list)  # redacted descriptions
    note: str = ""                                          # user-facing caveat, if any


# --- Secret scanning --------------------------------------------------------
# Best-effort, conservative patterns. The goal is to protect the user, not to be
# a full secret scanner. Each match is redacted in-place before anything is shown
# or sent to the model.
_SECRET_PATTERNS: tuple[tuple[str, re.Pattern[str]], None | str] = ()
SECRET_RULES: list[tuple[str, re.Pattern[str]]] = [
    ("OpenAI / Anthropic-style key", re.compile(r"\b(sk-[A-Za-z0-9_-]{16,}|sk-ant-[A-Za-z0-9_-]{16,})")),
    ("AWS access key", re.compile(r"\bAKIA[0-9A-Z]{16}\b")),
    ("Google API key", re.compile(r"\bAIza[0-9A-Za-z_-]{30,}\b")),
    ("GitHub token", re.compile(r"\bgh[pousr]_[A-Za-z0-9]{20,}\b")),
    ("Slack token", re.compile(r"\bxox[baprs]-[A-Za-z0-9-]{10,}\b")),
    ("Private key block", re.compile(r"-----BEGIN [A-Z ]*PRIVATE KEY-----")),
    # Generic "NAME = 'long-secret'" assignments for password/secret/token/api_key.
    ("Hard-coded secret value", re.compile(
        r"(?i)\b(password|passwd|secret|token|api[_-]?key|access[_-]?key)\b\s*[:=]\s*"
        r"['\"][^'\"]{8,}['\"]")),
]
_REDACTION = "«REDACTED-SECRET»"


def redact_secrets(text: str) -> tuple[str, list[str]]:
    """Replace likely secrets with a placeholder. Returns (clean_text, hits)."""
    hits: list[str] = []
    for label, pattern in SECRET_RULES:
        if pattern.search(text):
            hits.append(label)
            text = pattern.sub(_REDACTION, text)
    return text, hits


# --- Filtering --------------------------------------------------------------
def _lang_for(path: str) -> str | None:
    _, ext = os.path.splitext(path.lower())
    if ext in config.BINARY_EXTS:
        return None
    return config.LANG_BY_EXT.get(ext)


def _should_skip_dir(name: str) -> bool:
    return name in config.IGNORE_DIRS or name.startswith(".") or name == "__MACOSX"


def _iter_files(root: str):
    """Yield (relpath, abspath) for candidate files under root, pruning noise."""
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = [d for d in dirnames if not _should_skip_dir(d)]
        for fn in filenames:
            if fn.lower() in config.IGNORE_FILES:
                continue
            abspath = os.path.join(dirpath, fn)
            rel = os.path.relpath(abspath, root).replace(os.sep, "/")
            yield rel, abspath


def _read(abspath: str) -> str | None:
    try:
        if os.path.getsize(abspath) > config.MAX_FILE_BYTES:
            return None
        with open(abspath, "r", encoding="utf-8", errors="ignore") as fh:
            return fh.read()
    except (OSError, ValueError):
        return None


# --- Public entry points ----------------------------------------------------
def from_folder(root: str, name: str | None = None) -> Ingested:
    """Analyse a directory already on disk."""
    out = Ingested(name=name or os.path.basename(os.path.normpath(root)) or "your project")
    total = 0
    for rel, abspath in _iter_files(root):
        lang = _lang_for(rel)
        base = os.path.basename(rel).lower()
        is_manifest = base in config.DEP_MANIFESTS
        if lang is None and not is_manifest:
            out.skipped.append(rel)
            continue
        text = _read(abspath)
        if text is None:
            out.skipped.append(rel)
            continue
        total += len(text.encode("utf-8", "ignore"))
        if total > config.MAX_TOTAL_BYTES or len(out.files) >= config.MAX_FILES:
            out.skipped.append(rel)
            out.note = (f"Large project — analysed the first {len(out.files)} files; "
                        f"the rest are listed but not narrated.")
            continue
        clean, hits = redact_secrets(text)
        for h in hits:
            desc = f"{rel}: {h}"
            if desc not in out.secrets_found:
                out.secrets_found.append(desc)
        out.files.append(SourceFile(path=rel, abspath=abspath, text=clean,
                                    lang=lang or "other"))
    out.files.sort(key=lambda f: (0 if f.lang in ("python", "javascript", "typescript") else 1, f.path))
    out.skipped.sort()
    return out


def from_zip(zip_path: str, name: str | None = None) -> Ingested:
    """Unzip to a temp dir and analyse it."""
    tmp = tempfile.mkdtemp(prefix="storycode_")
    derived = name or os.path.splitext(os.path.basename(zip_path))[0]
    with zipfile.ZipFile(zip_path) as zf:
        for member in zf.namelist():
            # Guard against zip-slip; skip absolute / parent-escaping paths.
            target = os.path.normpath(os.path.join(tmp, member))
            if not target.startswith(os.path.abspath(tmp) + os.sep) and target != os.path.abspath(tmp):
                continue
            if member.endswith("/"):
                continue
            os.makedirs(os.path.dirname(target), exist_ok=True)
            try:
                with zf.open(member) as src, open(target, "wb") as dst:
                    dst.write(src.read())
            except (OSError, zipfile.BadZipFile):
                continue
    # A zip often wraps everything in a single top folder; descend into it.
    entries = [e for e in os.listdir(tmp) if not e.startswith("__MACOSX")]
    root = tmp
    if len(entries) == 1 and os.path.isdir(os.path.join(tmp, entries[0])):
        root = os.path.join(tmp, entries[0])
        derived = name or entries[0]
    return from_folder(root, name=derived)


def from_text(text: str, filename: str = "snippet.py") -> Ingested:
    """Analyse a single pasted file."""
    lang = _lang_for(filename) or "other"
    clean, hits = redact_secrets(text)
    out = Ingested(name=filename)
    out.files.append(SourceFile(path=filename, abspath="", text=clean, lang=lang))
    out.secrets_found = [f"{filename}: {h}" for h in hits]
    return out


# --- GitHub URL ingestion ---------------------------------------------------
_GH_URL_RE = re.compile(
    r"https?://github\.com/(?P<owner>[^/\s]+)/(?P<repo>[^/\s]+?)(?:\.git)?"
    r"(?:/(?:tree|blob)/(?P<branch>[^/\s]+))?"
    r"(?:/(?P<path>[^\s]*))?"
    r"\s*$"
)


def parse_github_url(url: str) -> tuple[str, str, str] | None:
    """Parse a GitHub URL into (owner, repo, branch). Returns None if invalid."""
    m = _GH_URL_RE.search(url.strip())
    if not m:
        return None
    owner = m.group("owner")
    repo = m.group("repo")
    branch = m.group("branch") or "HEAD"
    return owner, repo, branch


def from_github(url: str, name: str | None = None) -> Ingested:
    """Download a public GitHub repo and analyse it.

    Uses the GitHub archive endpoint (no git required). For large repos,
    only the first MAX_FILES files are analysed; the rest are listed but
    not narrated (handled by from_folder's size caps).
    """
    import urllib.request

    parsed = parse_github_url(url)
    if not parsed:
        return Ingested(name="invalid URL", note="Could not parse GitHub URL. Use: https://github.com/owner/repo")

    owner, repo, branch = parsed
    repo_name = name or f"{owner}/{repo}"

    # Try the archive endpoint: main → master fallback
    archive_urls = [
        f"https://github.com/{owner}/{repo}/archive/refs/heads/{branch}.zip",
        f"https://github.com/{owner}/{repo}/archive/refs/heads/main.zip",
        f"https://github.com/{owner}/{repo}/archive/refs/heads/master.zip",
    ]

    tmp = tempfile.mkdtemp(prefix="storycode_gh_")
    zip_path = os.path.join(tmp, "repo.zip")

    for archive_url in archive_urls:
        try:
            req = urllib.request.Request(archive_url, headers={"User-Agent": "StoryCode/1.0"})
            with urllib.request.urlopen(req, timeout=60) as resp:
                data = resp.read()
            with open(zip_path, "wb") as f:
                f.write(data)
            break
        except (urllib.error.URLError, OSError):
            continue
    else:
        return Ingested(
            name=repo_name,
            note="Could not download the repository. Make sure it's public and the URL is correct.",
        )

    # Extract and find the root
    with zipfile.ZipFile(zip_path) as zf:
        for member in zf.namelist():
            target = os.path.normpath(os.path.join(tmp, member))
            if not target.startswith(os.path.abspath(tmp) + os.sep) and target != os.path.abspath(tmp):
                continue
            if member.endswith("/"):
                continue
            os.makedirs(os.path.dirname(target), exist_ok=True)
            try:
                with zf.open(member) as src, open(target, "wb") as dst:
                    dst.write(src.read())
            except (OSError, zipfile.BadZipFile):
                continue

    # GitHub archives wrap in {repo}-{branch}/ folder
    entries = [e for e in os.listdir(tmp) if not e.startswith("__MACOSX") and e != "repo.zip"]
    root = tmp
    if len(entries) == 1 and os.path.isdir(os.path.join(tmp, entries[0])):
        root = os.path.join(tmp, entries[0])

    return from_folder(root, name=repo_name)