File size: 6,397 Bytes
6303ae6
 
 
 
 
 
 
 
 
 
 
 
 
357ee36
 
 
 
6303ae6
 
 
 
357ee36
 
6303ae6
 
 
 
 
 
 
 
 
 
357ee36
 
 
 
6303ae6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
357ee36
 
6303ae6
 
 
 
 
 
 
 
357ee36
 
 
 
 
 
 
 
6303ae6
 
357ee36
 
 
6303ae6
357ee36
 
6303ae6
 
357ee36
 
6303ae6
 
 
357ee36
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6303ae6
357ee36
 
 
 
 
 
 
 
6303ae6
 
357ee36
6303ae6
357ee36
 
 
 
6303ae6
 
 
357ee36
 
 
 
 
6303ae6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
"""File utilities used across services.

These helpers harden the dossier walk against the cases the project's
privacy and robustness guarantees depend on:

* **Symlink containment.** ``Path.is_file()`` follows symlinks, so a
  symlink dropped inside the chosen dossier folder (e.g.
  ``evidence.txt -> ~/.ssh/id_rsa``) would otherwise be read and its
  content shipped to the LLM, breaking "original files never leave the
  machine". :func:`iter_dossier_files` skips symlinked entries and any
  real path that resolves outside the chosen folder.
* **Size / count caps.** A single huge file or an enormous folder would
  otherwise be read fully into memory and hang the app. The walk is
  bounded by :data:`MAX_DOSSIER_FILES` and :data:`MAX_DOSSIER_WALK_ENTRIES`
  and individual oversized files are reported via
  :func:`is_within_size_limit` so callers can degrade gracefully instead
  of reading them.
"""

from __future__ import annotations

import os
from itertools import islice
from pathlib import Path
from typing import Iterator

from app.config import SUPPORTED_DOSSIER_EXTENSIONS


# Hard caps for the dossier walk. These bound a self-inflicted memory/time
# blow-up on a single-user local tool; they are deliberately generous.
MAX_DOSSIER_FILES = 2000
MAX_DOSSIER_FILE_BYTES = 25 * 1024 * 1024  # 25 MB per file
# Directory entries (files and folders) the walk reads in total, so a home
# directory or a project with node_modules stops the walk instead of
# hanging the app.
MAX_DOSSIER_WALK_ENTRIES = 20_000


def is_within_folder(path: Path, folder: Path) -> bool:
    """Return True if ``path`` resolves to a location inside ``folder``.

    Both paths are fully resolved (following symlinks) before comparison,
    so a symlink whose target escapes ``folder`` is rejected.
    """
    try:
        real = path.resolve()
        folder_real = folder.resolve()
    except OSError:
        return False
    if real == folder_real:
        return True
    return folder_real in real.parents


def is_within_size_limit(path: Path, max_bytes: int = MAX_DOSSIER_FILE_BYTES) -> bool:
    """Return True if ``path`` is at or below ``max_bytes``.

    Returns False for files whose size cannot be determined so an
    unreadable/odd entry is treated as "skip", not "read unbounded".
    """
    try:
        return path.stat().st_size <= max_bytes
    except OSError:
        return False


def iter_contained_files(
    folder: str | Path,
    *,
    max_files: int = MAX_DOSSIER_FILES,
    max_entries: int = MAX_DOSSIER_WALK_ENTRIES,
    stats: dict | None = None,
) -> Iterator[Path]:
    """Yield regular files under ``folder`` that are safe to consider.

    This is the security primitive the dossier walk is built on:

    * **Symlinks are skipped** — both symlinked files and any entry whose
      resolved real path falls outside ``folder`` (containment check), so
      content from elsewhere on disk can never enter the dossier.
    * **The walk is lazy and bounded** so an enormous folder cannot hang
      the app. Folders are read one at a time, each in name order (the
      order ``sorted(rglob("*"))`` gave), so files come out without the
      whole tree being listed first. At most ``max_entries`` directory
      entries are read, and the walk stops after ``max_files`` supported
      dossier files. Callers are responsible for surfacing truncation: pass
      a ``stats`` dict and ``stats["truncated"]`` says whether either limit
      cut the walk short.

    No extension filtering is applied here, so callers that need to report
    unsupported files (e.g. the folder validator) still see them. They do
    not count toward ``max_files``, so a folder full of images, code or
    build output cannot crowd out the real dossier files.

    Symlinked *directories* are never descended into, so a directory-symlink
    loop is not reachable through this walk.
    """
    folder = Path(folder).expanduser()
    if stats is not None:
        stats["truncated"] = False
    if not folder.exists() or not folder.is_dir():
        return

    entries_left = max_entries

    def truncate() -> None:
        if stats is not None:
            stats["truncated"] = True

    def sorted_entries(directory: Path) -> Iterator[os.DirEntry]:
        # Read no more than the entry budget allows (one extra tells us
        # whether anything was left unread), then sort by name.
        nonlocal entries_left
        try:
            with os.scandir(directory) as it:
                entries = list(islice(it, entries_left + 1))
        except OSError:
            # An unreadable sub-folder is skipped, as rglob did.
            return iter(())
        if len(entries) > entries_left:
            entries = entries[:entries_left]
            truncate()
        entries_left -= len(entries)
        return iter(sorted(entries, key=lambda entry: entry.name))

    yielded = 0
    stack = [(folder, sorted_entries(folder))]
    while stack:
        directory, entries = stack[-1]
        entry = next(entries, None)
        if entry is None:
            stack.pop()
            continue
        path = directory / entry.name
        # Skip symlinks outright (both files and dirs) and anything that
        # resolves outside the chosen folder.
        if entry.is_symlink():
            continue
        if entry.is_dir(follow_symlinks=False):
            stack.append((path, sorted_entries(path)))
            continue
        if not entry.is_file(follow_symlinks=False):
            continue
        if not is_within_folder(path, folder):
            continue
        if path.suffix.lower() in SUPPORTED_DOSSIER_EXTENSIONS:
            if yielded >= max_files:
                truncate()
                return
            yielded += 1
        yield path


def iter_dossier_files(
    folder: str | Path,
    *,
    max_files: int = MAX_DOSSIER_FILES,
) -> Iterator[Path]:
    """Yield contained, supported-extension dossier files under ``folder``."""
    for path in iter_contained_files(folder, max_files=max_files):
        if path.suffix.lower() in SUPPORTED_DOSSIER_EXTENSIONS:
            yield path


def iter_supported_files(folder: str | Path):
    """Backwards-compatible alias for :func:`iter_dossier_files`.

    Retained so existing imports keep working; new code should call
    :func:`iter_dossier_files` directly.
    """
    yield from iter_dossier_files(folder)