Skills › Scientific Agent Skills › exploratory-data-analysis › _common.py
_common.py
468 lines · skills/exploratory-data-analysis/scripts/_common.py at 5f2bc03
1#!/usr/bin/env python3
2"""Shared local-only, bounded-I/O helpers for the EDA command-line tools."""
4from __future__ import annotations
6import hashlib
7import json
8import math
9import os
10import stat
11import tempfile
12from collections.abc import Iterable, Mapping
13from pathlib import Path, PurePath
14from typing import Any
17SCHEMA_VERSION = "1.1"
18MIB = 1024 * 1024
19DEFAULT_MAX_FILE_BYTES = 64 * MIB
20MAX_FILE_BYTES = 512 * MIB
21MAX_JSON_BYTES = 16 * MIB
22MAX_REPORT_BYTES = 4 * MIB
23MAX_ROWS = 1_000_000
24DEFAULT_MAX_ROWS = 100_000
25MAX_COLUMNS = 512
26MAX_FIELD_CHARS = 100_000
27MAX_IDENTIFIER_CHARS = 160
28MAX_NPZ_MEMBERS = 128
29MAX_NPZ_UNCOMPRESSED_BYTES = 128 * MIB
30MAX_COMPRESSION_RATIO = 100.0
31MAX_IMAGE_PIXELS = 100_000_000
34class CliError(ValueError):
35 """An expected command-line validation error with no sensitive values."""
38def bounded_integer(
39 value: int,
40 *,
41 name: str,
42 minimum: int,
43 maximum: int,
44) -> int:
45 """Validate a non-boolean integer against fixed limits."""
47 if isinstance(value, bool) or not isinstance(value, int):
48 raise CliError(f"{name} must be an integer")
49 if not minimum <= value <= maximum:
50 raise CliError(f"{name} must be between {minimum} and {maximum}")
51 return value
54def bounded_file_limit(value: int) -> int:
55 """Validate a caller-selected byte limit against the hard ceiling."""
57 return bounded_integer(
58 value,
59 name="max bytes",
60 minimum=1,
61 maximum=MAX_FILE_BYTES,
62 )
65def finite_number(value: str) -> float | None:
66 """Parse a finite decimal number without evaluating expressions."""
68 try:
69 number = float(value)
70 except (TypeError, ValueError):
71 return None
72 return number if math.isfinite(number) else None
75def _reject_url_and_traversal(value: str) -> None:
76 """Reject URLs, NUL bytes, home expansion, and lexical parent traversal."""
78 stripped = value.strip()
79 lowered = stripped.lower()
80 if not stripped:
81 raise CliError("path must not be empty")
82 if "\x00" in value:
83 raise CliError("path must not contain a NUL byte")
84 if "://" in lowered or lowered.startswith(
85 ("http:", "https:", "ftp:", "s3:", "gs:", "file:", "data:")
86 ):
87 raise CliError("URLs are not accepted; provide a local path")
88 if stripped.startswith("~"):
89 raise CliError("home-directory expansion is not accepted")
90 if ".." in PurePath(stripped).parts:
91 raise CliError("parent-directory traversal is not accepted")
94def _absolute_lexical(path: Path) -> Path:
95 """Make a path absolute without intentionally resolving symlinks."""
97 return Path(os.path.abspath(os.fspath(path)))
100def _reject_symlink_components(path: Path) -> None:
101 """Reject an existing symlink at any component of an absolute path."""
103 absolute = _absolute_lexical(path)
104 current = Path(absolute.anchor)
105 for part in absolute.parts[1:]:
106 current /= part
107 try:
108 if current.is_symlink():
109 raise CliError("symlink path components are not accepted")
110 except OSError as exc:
111 raise CliError("a path component could not be inspected") from exc
114def checked_root(value: str | os.PathLike[str]) -> Path:
115 """Return an existing, non-symlink directory used as the I/O boundary."""
117 raw = os.fspath(value)
118 _reject_url_and_traversal(raw)
119 supplied = _absolute_lexical(Path(raw))
120 if supplied.is_symlink():
121 raise CliError("the root directory must not be a symlink")
122 try:
123 root = supplied.resolve(strict=True)
124 info = root.stat()
125 except OSError as exc:
126 raise CliError("the root directory is not accessible") from exc
127 if not stat.S_ISDIR(info.st_mode):
128 raise CliError("the root must be an existing directory")
129 _reject_symlink_components(root)
130 return root
133def _within_root(candidate: Path, root: Path) -> None:
134 try:
135 candidate.relative_to(root)
136 except ValueError as exc:
137 raise CliError("path escapes the declared root directory") from exc
140def _suffix_matches(path: Path, suffixes: Iterable[str]) -> bool:
141 name = path.name.casefold()
142 return any(name.endswith(suffix.casefold()) for suffix in suffixes)
145def checked_input_file(
146 value: str | os.PathLike[str],
147 *,
148 root: str | os.PathLike[str] = ".",
149 suffixes: Iterable[str] | None = None,
150 max_bytes: int = DEFAULT_MAX_FILE_BYTES,
151) -> Path:
152 """Return a bounded regular local file inside root, rejecting all symlinks."""
154 max_bytes = bounded_file_limit(max_bytes)
155 raw = os.fspath(value)
156 _reject_url_and_traversal(raw)
157 root_path = checked_root(root)
158 path = Path(raw)
159 if not path.is_absolute():
160 path = root_path / path
161 path = _absolute_lexical(path)
162 if path.is_symlink():
163 raise CliError("the input must not be a symlink")
164 try:
165 resolved = path.resolve(strict=True)
166 info = resolved.stat()
167 except OSError as exc:
168 raise CliError("the input file is not accessible") from exc
169 _within_root(resolved, root_path)
170 _reject_symlink_components(resolved)
171 if not stat.S_ISREG(info.st_mode):
172 raise CliError("the input must be a regular file")
173 if info.st_nlink != 1:
174 raise CliError("multiply linked input files are not accepted")
175 if info.st_size > max_bytes:
176 raise CliError(
177 f"the input is {info.st_size} bytes; the configured limit is {max_bytes}"
178 )
179 if suffixes is not None and not _suffix_matches(resolved, suffixes):
180 allowed = ", ".join(sorted({item.casefold() for item in suffixes}))
181 raise CliError(f"the input suffix must be one of: {allowed}")
182 return resolved
185def checked_output_file(
186 value: str | os.PathLike[str],
187 *,
188 root: str | os.PathLike[str] = ".",
189 suffixes: Iterable[str],
190 force: bool = False,
191) -> Path:
192 """Return a local output path inside root without following symlinks."""
194 raw = os.fspath(value)
195 _reject_url_and_traversal(raw)
196 root_path = checked_root(root)
197 path = Path(raw)
198 if not path.is_absolute():
199 path = root_path / path
200 path = _absolute_lexical(path)
201 if path.name in {"", ".", ".."}:
202 raise CliError("the output must name a file")
203 if not _suffix_matches(path, suffixes):
204 allowed = ", ".join(sorted({item.casefold() for item in suffixes}))
205 raise CliError(f"the output suffix must be one of: {allowed}")
206 if path.is_symlink() or path.parent.is_symlink():
207 raise CliError("output symlinks are not accepted")
208 try:
209 parent = path.parent.resolve(strict=True)
210 parent_info = parent.stat()
211 except OSError as exc:
212 raise CliError("the output parent is not accessible") from exc
213 _within_root(parent, root_path)
214 _reject_symlink_components(parent)
215 if not stat.S_ISDIR(parent_info.st_mode):
216 raise CliError("the output parent must be an existing directory")
217 destination = parent / path.name
218 if destination.exists():
219 if destination.is_symlink() or not destination.is_file():
220 raise CliError("the output exists but is not a regular file")
221 if not force:
222 raise CliError("refusing to overwrite an existing output")
223 return destination
226def _strict_json_bytes(document: Any) -> bytes:
227 try:
228 payload = (
229 json.dumps(
230 document,
231 allow_nan=False,
232 ensure_ascii=False,
233 indent=2,
234 sort_keys=True,
235 )
236 + "\n"
237 ).encode("utf-8")
238 except (TypeError, ValueError) as exc:
239 raise CliError("the report could not be serialized as strict JSON") from exc
240 if len(payload) > MAX_REPORT_BYTES:
241 raise CliError(
242 f"the report is {len(payload)} bytes; the limit is {MAX_REPORT_BYTES}"
243 )
244 return payload
247def atomic_write_bytes(
248 output: str | os.PathLike[str],
249 payload: bytes,
250 *,
251 root: str | os.PathLike[str] = ".",
252 suffixes: Iterable[str],
253 force: bool = False,
254) -> Path:
255 """Write a private file atomically in an existing local directory."""
257 destination = checked_output_file(
258 output,
259 root=root,
260 suffixes=suffixes,
261 force=force,
262 )
263 descriptor, temporary_name = tempfile.mkstemp(
264 prefix=f".{destination.name}.",
265 suffix=".tmp",
266 dir=destination.parent,
267 )
268 temporary = Path(temporary_name)
269 try:
270 with os.fdopen(descriptor, "wb") as handle:
271 handle.write(payload)
272 handle.flush()
273 os.fsync(handle.fileno())
274 os.chmod(temporary, 0o600)
275 if destination.exists() and not force:
276 raise CliError("refusing to overwrite an existing output")
277 os.replace(temporary, destination)
278 finally:
279 temporary.unlink(missing_ok=True)
280 return destination
283def emit_json(
284 document: Any,
285 *,
286 output: str | os.PathLike[str] | None = None,
287 root: str | os.PathLike[str] = ".",
288 force: bool = False,
289) -> None:
290 """Print strict JSON or write it privately and atomically."""
292 payload = _strict_json_bytes(document)
293 if output is None:
294 print(payload.decode("utf-8"), end="")
295 return
296 atomic_write_bytes(
297 output,
298 payload,
299 root=root,
300 suffixes={".json"},
301 force=force,
302 )
305def emit_markdown(
306 text: str,
307 *,
308 output: str | os.PathLike[str] | None = None,
309 root: str | os.PathLike[str] = ".",
310 force: bool = False,
311) -> None:
312 """Print bounded Markdown or write it privately and atomically."""
314 payload = text.encode("utf-8")
315 if len(payload) > MAX_REPORT_BYTES:
316 raise CliError(
317 f"the report is {len(payload)} bytes; the limit is {MAX_REPORT_BYTES}"
318 )
319 if output is None:
320 print(text, end="" if text.endswith("\n") else "\n")
321 return
322 atomic_write_bytes(
323 output,
324 payload,
325 root=root,
326 suffixes={".md", ".markdown"},
327 force=force,
328 )
331def stable_token(value: str, *, kind: str) -> str:
332 """Return a deterministic pseudonymous token; this is not anonymization."""
334 digest = hashlib.blake2s(
335 f"eda-v1.1\0{kind}\0{value}".encode("utf-8", errors="surrogatepass"),
336 digest_size=8,
337 ).hexdigest()
338 return f"{kind}_{digest}"
341def sanitize_identifier(value: str) -> str:
342 """Bound and neutralize an explicitly requested untrusted identifier."""
344 cleaned = "".join(
345 character if character.isprintable() and character not in "\r\n\t" else " "
346 for character in value
347 )
348 cleaned = " ".join(cleaned.split())
349 if len(cleaned) > MAX_IDENTIFIER_CHARS:
350 cleaned = cleaned[:MAX_IDENTIFIER_CHARS] + "…"
351 return cleaned
354def markdown_scalar(value: str) -> str:
355 """Return a bounded scalar that cannot introduce Markdown structure."""
357 cleaned = sanitize_identifier(value)
358 allowed_punctuation = frozenset(" .,_-:()/+")
359 return "".join(
360 character
361 if character.isalnum() or character in allowed_punctuation
362 else "�"
363 for character in cleaned
364 )
367def display_identifier(
368 value: str,
369 *,
370 kind: str,
371 reveal_identifiers: bool,
372) -> str:
373 """Reveal a sanitized identifier only after explicit caller opt-in."""
375 if reveal_identifiers:
376 return sanitize_identifier(value)
377 return stable_token(value, kind=kind)
380def _reject_json_constant(value: str) -> None:
381 raise CliError("non-finite JSON numbers are not accepted")
384def _finite_json_float(value: str) -> float:
385 number = float(value)
386 if not math.isfinite(number):
387 raise CliError("JSON number exceeds the finite floating-point range")
388 return number
391def _unique_json_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
392 result: dict[str, Any] = {}
393 for key, value in pairs:
394 if key in result:
395 raise CliError("duplicate JSON object keys are not accepted")
396 result[key] = value
397 return result
400def load_strict_json(
401 path: Path,
402 *,
403 max_bytes: int = MAX_JSON_BYTES,
404) -> Any:
405 """Load bounded RFC-style JSON, rejecting duplicate keys and NaN/Infinity."""
407 if path.stat().st_size > max_bytes:
408 raise CliError(f"JSON input exceeds the {max_bytes}-byte parsing limit")
409 try:
410 with path.open("r", encoding="utf-8") as handle:
411 return json.load(
412 handle,
413 object_pairs_hook=_unique_json_object,
414 parse_constant=_reject_json_constant,
415 parse_float=_finite_json_float,
416 )
417 except CliError:
418 raise
419 except (OSError, UnicodeError, ValueError, RecursionError, MemoryError) as exc:
420 raise CliError("the input is not bounded, valid UTF-8 JSON") from exc
423def validate_keys(
424 value: Mapping[str, Any],
425 *,
426 allowed: Iterable[str],
427 required: Iterable[str] = (),
428 context: str,
429) -> None:
430 """Reject unknown keys and report missing required keys."""
432 allowed_set = set(allowed)
433 required_set = set(required)
434 unknown = set(value) - allowed_set
435 missing = required_set - set(value)
436 if unknown:
437 raise CliError(f"{context} contains unsupported keys")
438 if missing:
439 raise CliError(f"{context} is missing required keys")
442def sha256_file(path: Path, *, max_bytes: int) -> str:
443 """Hash a previously checked bounded regular file using streaming reads."""
445 if path.stat().st_size > max_bytes:
446 raise CliError("the file exceeds the configured hashing limit")
447 digest = hashlib.sha256()
448 try:
449 with path.open("rb") as handle:
450 for chunk in iter(lambda: handle.read(MIB), b""):
451 digest.update(chunk)
452 except OSError as exc:
453 raise CliError("the input could not be hashed") from exc
454 return digest.hexdigest()
457def run_cli(function: Any) -> int:
458 """Run a CLI body with concise expected-error handling."""
460 try:
461 function()
462 except CliError as exc:
463 print(f"error: {exc}", file=os.sys.stderr)
464 return 2
465 except KeyboardInterrupt:
466 print("error: interrupted", file=os.sys.stderr)
467 return 130
468 return 0