Skills › Scientific Agent Skills › fluidsim › _common.py
_common.py
494 lines · skills/fluidsim/scripts/_common.py at 5f2bc03
1#!/usr/bin/env python3
2"""Shared dependency-free safety and JSON helpers for FluidSim skill CLIs."""
4from __future__ import annotations
6import hashlib
7import json
8import math
9import os
10import re
11import stat
12import sys
13import tempfile
14from collections.abc import Iterable, Mapping
15from pathlib import Path, PurePath
16from typing import Any
19SCHEMA_VERSION = "1.1"
20MIB = 1024**2
21GIB = 1024**3
22MAX_JSON_BYTES = 4 * MIB
23MAX_REPORT_BYTES = 8 * MIB
24MAX_TEXT_BYTES = 256 * MIB
25MAX_HDF5_BYTES = 8 * 1024 * GIB
26MAX_FILES = 10_000
27MAX_DATASETS = 100_000
28MAX_ATTRIBUTES = 100_000
29MAX_RECORDS = 2_000_000
30MAX_GRID_POINTS = 2**48
31MAX_DIMENSION = 131_072
32MAX_OUTPUT_FILES = 1_000_000
33MAX_CPU_CORES = 1_048_576
34MAX_WALL_MINUTES = 10 * 365 * 24 * 60
36_URI_PREFIXES = (
37 "http:",
38 "https:",
39 "ftp:",
40 "file:",
41 "s3:",
42 "gs:",
43 "ssh:",
44 "data:",
45)
46_SAFE_SLUG = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$")
49class ToolError(ValueError):
50 """Expected validation or local-I/O failure."""
53def bounded_int(
54 value: Any, *, name: str, minimum: int, maximum: int
55) -> int:
56 """Validate a non-boolean integer against fixed limits."""
58 if isinstance(value, bool) or not isinstance(value, int):
59 raise ToolError(f"{name} must be an integer")
60 if not minimum <= value <= maximum:
61 raise ToolError(f"{name} must be between {minimum} and {maximum}")
62 return value
65def finite_float(
66 value: Any,
67 *,
68 name: str,
69 minimum: float | None = None,
70 maximum: float | None = None,
71 allow_none: bool = False,
72) -> float | None:
73 """Validate a finite JSON number."""
75 if value is None and allow_none:
76 return None
77 if isinstance(value, bool) or not isinstance(value, (int, float)):
78 raise ToolError(f"{name} must be a finite number")
79 try:
80 number = float(value)
81 except OverflowError as exc:
82 raise ToolError(f"{name} must be a finite number") from exc
83 if not math.isfinite(number):
84 raise ToolError(f"{name} must be a finite number")
85 if minimum is not None and number < minimum:
86 raise ToolError(f"{name} must be at least {minimum}")
87 if maximum is not None and number > maximum:
88 raise ToolError(f"{name} must be at most {maximum}")
89 return number
92def validate_keys(
93 value: Mapping[str, Any],
94 *,
95 allowed: Iterable[str],
96 required: Iterable[str] = (),
97 context: str,
98) -> None:
99 """Reject unknown keys and require declared keys."""
101 allowed_set = set(allowed)
102 required_set = set(required)
103 unknown = sorted(set(value) - allowed_set)
104 missing = sorted(required_set - set(value))
105 if unknown:
106 raise ToolError(f"{context} contains unsupported keys: {', '.join(unknown)}")
107 if missing:
108 raise ToolError(f"{context} is missing keys: {', '.join(missing)}")
111def require_mapping(value: Any, *, context: str) -> Mapping[str, Any]:
112 if not isinstance(value, Mapping):
113 raise ToolError(f"{context} must be a JSON object")
114 return value
117def require_text(
118 value: Any, *, name: str, minimum: int = 1, maximum: int = 2_000
119) -> str:
120 if not isinstance(value, str):
121 raise ToolError(f"{name} must be a string")
122 if not minimum <= len(value) <= maximum:
123 raise ToolError(f"{name} length must be between {minimum} and {maximum}")
124 if any(ord(character) < 32 and character not in "\t" for character in value):
125 raise ToolError(f"{name} contains control characters")
126 return value
129def require_text_list(
130 value: Any, *, name: str, minimum: int = 1, maximum: int = 100
131) -> list[str]:
132 if not isinstance(value, list) or not minimum <= len(value) <= maximum:
133 raise ToolError(f"{name} must contain {minimum}..{maximum} strings")
134 return [
135 require_text(item, name=f"{name} item", minimum=1, maximum=1_000)
136 for item in value
137 ]
140def safe_slug(value: Any, *, name: str) -> str:
141 """Validate one local path component."""
143 text = require_text(value, name=name, maximum=128)
144 if not _SAFE_SLUG.fullmatch(text) or text in {".", ".."}:
145 raise ToolError(
146 f"{name} must use only letters, digits, '.', '_', or '-'"
147 )
148 return text
151def safe_relative_path(
152 value: Any,
153 *,
154 name: str,
155 suffixes: Iterable[str] | None = None,
156 allow_nested: bool = False,
157) -> str:
158 """Validate a local relative path without traversal or URI syntax."""
160 text = require_text(value, name=name, maximum=512).strip()
161 lowered = text.casefold()
162 if (
163 text.startswith(("~", "/", "\\"))
164 or "://" in lowered
165 or lowered.startswith(_URI_PREFIXES)
166 ):
167 raise ToolError(f"{name} must be a local relative path")
168 path = PurePath(text)
169 if path.is_absolute() or ".." in path.parts or "." in path.parts:
170 raise ToolError(f"{name} must not contain traversal")
171 if not allow_nested and len(path.parts) != 1:
172 raise ToolError(f"{name} must be one path component")
173 for part in path.parts:
174 safe_slug(part, name=name)
175 if suffixes is not None and not any(
176 text.casefold().endswith(suffix.casefold()) for suffix in suffixes
177 ):
178 raise ToolError(f"{name} has an unsupported suffix")
179 return text
182def _reject_unsafe_path_text(value: str) -> None:
183 stripped = value.strip()
184 lowered = stripped.casefold()
185 if not stripped or "\x00" in value:
186 raise ToolError("path must be nonempty and contain no NUL byte")
187 if (
188 stripped.startswith("~")
189 or "://" in lowered
190 or lowered.startswith(_URI_PREFIXES)
191 ):
192 raise ToolError("only local filesystem paths are accepted")
193 if ".." in PurePath(stripped).parts:
194 raise ToolError("parent traversal is not accepted")
197def _absolute_lexical(path: Path) -> Path:
198 return Path(os.path.abspath(os.fspath(path)))
201def _reject_symlink_components(path: Path) -> None:
202 absolute = _absolute_lexical(path)
203 current = Path(absolute.anchor)
204 for part in absolute.parts[1:]:
205 current /= part
206 try:
207 mode = current.lstat().st_mode
208 except FileNotFoundError:
209 continue
210 except OSError as exc:
211 raise ToolError("a path component could not be inspected") from exc
212 if stat.S_ISLNK(mode):
213 raise ToolError("symlink path components are not accepted")
216def checked_root(value: str | os.PathLike[str]) -> Path:
217 """Resolve an existing, non-symlink local directory."""
219 raw = os.fspath(value)
220 _reject_unsafe_path_text(raw)
221 supplied = _absolute_lexical(Path(raw))
222 _reject_symlink_components(supplied)
223 try:
224 resolved = supplied.resolve(strict=True)
225 info = resolved.stat()
226 except OSError as exc:
227 raise ToolError("root directory is not accessible") from exc
228 if not stat.S_ISDIR(info.st_mode):
229 raise ToolError("root must be a directory")
230 return resolved
233def _within_root(path: Path, root: Path) -> None:
234 try:
235 path.relative_to(root)
236 except ValueError as exc:
237 raise ToolError("path escapes the declared root") from exc
240def checked_input(
241 value: str | os.PathLike[str],
242 *,
243 root: str | os.PathLike[str] = ".",
244 kind: str = "file",
245 suffixes: Iterable[str] | None = None,
246 max_bytes: int = MAX_HDF5_BYTES,
247) -> Path:
248 """Resolve a bounded local file or directory without following symlinks."""
250 raw = os.fspath(value)
251 _reject_unsafe_path_text(raw)
252 root_path = checked_root(root)
253 candidate = Path(raw)
254 if not candidate.is_absolute():
255 candidate = root_path / candidate
256 candidate = _absolute_lexical(candidate)
257 _reject_symlink_components(candidate)
258 try:
259 resolved = candidate.resolve(strict=True)
260 info = resolved.stat()
261 except OSError as exc:
262 raise ToolError("input path is not accessible") from exc
263 _within_root(resolved, root_path)
264 if kind == "file" and not stat.S_ISREG(info.st_mode):
265 raise ToolError("input must be a regular file")
266 if kind == "dir" and not stat.S_ISDIR(info.st_mode):
267 raise ToolError("input must be a directory")
268 if kind == "any" and not (
269 stat.S_ISREG(info.st_mode) or stat.S_ISDIR(info.st_mode)
270 ):
271 raise ToolError("input must be a regular file or directory")
272 if stat.S_ISREG(info.st_mode):
273 if info.st_nlink != 1:
274 raise ToolError("multiply linked input files are not accepted")
275 if not 0 <= info.st_size <= max_bytes:
276 raise ToolError("input file exceeds the configured byte limit")
277 if suffixes is not None and not any(
278 resolved.name.casefold().endswith(suffix.casefold())
279 for suffix in suffixes
280 ):
281 raise ToolError("input file has an unsupported suffix")
282 return resolved
285def checked_output(
286 value: str | os.PathLike[str],
287 *,
288 root: str | os.PathLike[str] = ".",
289 suffixes: Iterable[str],
290 force: bool = False,
291) -> Path:
292 """Resolve a safe local output file within root."""
294 relative = safe_relative_path(
295 os.fspath(value),
296 name="output",
297 suffixes=suffixes,
298 allow_nested=True,
299 )
300 root_path = checked_root(root)
301 candidate = _absolute_lexical(root_path / relative)
302 parent = candidate.parent
303 _reject_symlink_components(parent)
304 try:
305 parent = parent.resolve(strict=True)
306 except OSError as exc:
307 raise ToolError("output parent must already exist") from exc
308 _within_root(parent, root_path)
309 destination = parent / candidate.name
310 if destination.exists():
311 if destination.is_symlink() or not destination.is_file():
312 raise ToolError("existing output is not a regular file")
313 if not force:
314 raise ToolError("refusing to overwrite existing output")
315 return destination
318def atomic_write(
319 destination: Path, payload: bytes, *, force: bool = False
320) -> None:
321 """Write a private file atomically without creating parent directories."""
323 if len(payload) > MAX_REPORT_BYTES:
324 raise ToolError("generated output exceeds the hard report-size limit")
325 descriptor, temporary_name = tempfile.mkstemp(
326 prefix=f".{destination.name}.", suffix=".tmp", dir=destination.parent
327 )
328 temporary = Path(temporary_name)
329 try:
330 with os.fdopen(descriptor, "wb") as handle:
331 handle.write(payload)
332 handle.flush()
333 os.fsync(handle.fileno())
334 os.chmod(temporary, 0o600)
335 if destination.exists() and not force:
336 raise ToolError("refusing to overwrite existing output")
337 os.replace(temporary, destination)
338 finally:
339 temporary.unlink(missing_ok=True)
342def _reject_json_constant(value: str) -> None:
343 raise ToolError(f"non-finite JSON number is not accepted: {value}")
346def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
347 result: dict[str, Any] = {}
348 for key, value in pairs:
349 if key in result:
350 raise ToolError(f"duplicate JSON key is not accepted: {key}")
351 result[key] = value
352 return result
355def load_json(path: Path, *, max_bytes: int = MAX_JSON_BYTES) -> Any:
356 """Load bounded strict UTF-8 JSON."""
358 if path.stat().st_size > max_bytes:
359 raise ToolError("JSON input exceeds the parsing limit")
360 try:
361 with path.open("r", encoding="utf-8") as handle:
362 return json.load(
363 handle,
364 object_pairs_hook=_unique_object,
365 parse_constant=_reject_json_constant,
366 )
367 except ToolError:
368 raise
369 except (OSError, UnicodeError, json.JSONDecodeError, RecursionError) as exc:
370 raise ToolError("input is not bounded, valid UTF-8 JSON") from exc
373def strict_json_loads(text: str) -> Any:
374 try:
375 return json.loads(
376 text,
377 object_pairs_hook=_unique_object,
378 parse_constant=_reject_json_constant,
379 )
380 except ToolError:
381 raise
382 except (json.JSONDecodeError, RecursionError) as exc:
383 raise ToolError("record is not strict JSON") from exc
386def json_bytes(document: Any) -> bytes:
387 try:
388 payload = (
389 json.dumps(
390 document,
391 allow_nan=False,
392 ensure_ascii=False,
393 indent=2,
394 sort_keys=True,
395 )
396 + "\n"
397 ).encode("utf-8")
398 except (TypeError, ValueError) as exc:
399 raise ToolError("report cannot be serialized as strict JSON") from exc
400 if len(payload) > MAX_REPORT_BYTES:
401 raise ToolError("report exceeds the hard report-size limit")
402 return payload
405def emit_json(document: Any) -> None:
406 sys.stdout.buffer.write(json_bytes(document))
409def fail_json(tool: str, exc: Exception) -> int:
410 emit_json(
411 {
412 "error": type(exc).__name__,
413 "message": str(exc)[:500],
414 "ok": False,
415 "tool": tool,
416 }
417 )
418 return 2
421def sha256_file(path: Path, *, max_bytes: int) -> str | None:
422 """Hash a regular file when it is within an explicit bound."""
424 size = path.stat().st_size
425 if size > max_bytes:
426 return None
427 digest = hashlib.sha256()
428 with path.open("rb") as handle:
429 for chunk in iter(lambda: handle.read(MIB), b""):
430 digest.update(chunk)
431 return digest.hexdigest()
434def iter_local_files(
435 path: Path,
436 *,
437 suffixes: Iterable[str] | None,
438 max_files: int,
439 recursive: bool = True,
440) -> list[Path]:
441 """List bounded regular files without traversing symlinks."""
443 bounded_int(max_files, name="max_files", minimum=1, maximum=MAX_FILES)
444 if path.is_file():
445 candidates = [path]
446 else:
447 candidates = []
448 stack = [path]
449 while stack:
450 directory = stack.pop()
451 try:
452 entries = sorted(
453 os.scandir(directory), key=lambda entry: entry.name.casefold()
454 )
455 except OSError as exc:
456 raise ToolError("directory cannot be scanned") from exc
457 for entry in entries:
458 try:
459 if entry.is_symlink():
460 raise ToolError("symlinks in scanned directories are rejected")
461 if entry.is_dir(follow_symlinks=False):
462 if recursive:
463 stack.append(Path(entry.path))
464 elif entry.is_file(follow_symlinks=False):
465 file_path = Path(entry.path)
466 info = file_path.stat()
467 if info.st_nlink != 1:
468 raise ToolError(
469 "multiply linked files in scanned directories are rejected"
470 )
471 candidates.append(file_path)
472 if len(candidates) > max_files:
473 raise ToolError("file-count limit exceeded")
474 except OSError as exc:
475 raise ToolError("directory entry cannot be inspected") from exc
476 if suffixes is None:
477 return candidates
478 lowered = tuple(suffix.casefold() for suffix in suffixes)
479 return [
480 candidate
481 for candidate in candidates
482 if candidate.name.casefold().endswith(lowered)
483 ]
486def relative_display(path: Path, root: Path) -> str:
487 """Return a bounded path relative to the caller-declared root."""
489 try:
490 relative = path.relative_to(root)
491 except ValueError:
492 return "<outside-root>"
493 text = relative.as_posix()
494 return text[:512]