Skills › Scientific Agent Skills › get-available-resources › _common.py
_common.py
199 lines · skills/get-available-resources/scripts/_common.py at 5f2bc03
1#!/usr/bin/env python3
2"""Shared, dependency-free safety and JSON helpers."""
4from __future__ import annotations
6import argparse
7import json
8import math
9import os
10import stat
11import sys
12from pathlib import Path
13from typing import Any
15SCHEMA_VERSION = "1.1"
16MAX_SNAPSHOT_BYTES = 1_048_576
17MAX_WORKERS = 1_024
18MAX_TASKS = 1_000_000_000
19MAX_BYTES = 1 << 63
22class ResourceToolError(ValueError):
23 """A concise, user-safe CLI error."""
26def json_text(value: Any) -> str:
27 """Serialize deterministic, standards-compliant JSON."""
28 try:
29 return (
30 json.dumps(
31 value,
32 allow_nan=False,
33 ensure_ascii=False,
34 indent=2,
35 sort_keys=True,
36 )
37 + "\n"
38 )
39 except (RecursionError, TypeError, ValueError) as exc:
40 raise ResourceToolError("result is not valid JSON data") from exc
43def _safe_output_path(raw: str) -> Path:
44 """Restrict writes to one explicit JSON filename in the current directory."""
45 if not raw or len(raw) > 255:
46 raise ResourceToolError("output must be a short JSON filename")
47 candidate = Path(raw)
48 if candidate.is_absolute() or candidate.drive or len(candidate.parts) != 1:
49 raise ResourceToolError(
50 "output must be a filename in the current directory, not a path"
51 )
52 if candidate.name in {".", ".."} or candidate.suffix.lower() != ".json":
53 raise ResourceToolError("output filename must end in .json")
54 return Path.cwd() / candidate.name
57def emit_json(value: Any, output: str | None = None, *, force: bool = False) -> None:
58 """Write JSON to stdout or an explicit private local file."""
59 payload = json_text(value)
60 if output is None:
61 sys.stdout.write(payload)
62 return
64 destination = _safe_output_path(output)
65 try:
66 if destination.is_symlink():
67 raise ResourceToolError("refusing to write through a symbolic link")
68 flags = os.O_WRONLY | os.O_CREAT
69 if not force:
70 flags |= os.O_EXCL
71 flags |= getattr(os, "O_NONBLOCK", 0)
72 if hasattr(os, "O_NOFOLLOW"):
73 flags |= os.O_NOFOLLOW
74 descriptor = os.open(destination, flags, 0o600)
75 try:
76 info = os.fstat(descriptor)
77 if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1:
78 raise ResourceToolError("output must be a regular file without hard links")
79 if force:
80 os.ftruncate(descriptor, 0)
81 if hasattr(os, "fchmod"):
82 os.fchmod(descriptor, 0o600)
83 with os.fdopen(descriptor, "w", encoding="utf-8") as stream:
84 descriptor = -1
85 stream.write(payload)
86 stream.flush()
87 os.fsync(stream.fileno())
88 finally:
89 if descriptor >= 0:
90 os.close(descriptor)
91 except FileExistsError as exc:
92 raise ResourceToolError(
93 f"{destination.name} already exists; use --force to replace it"
94 ) from exc
95 except OSError as exc:
96 raise ResourceToolError(
97 f"could not write the requested output file ({exc.__class__.__name__})"
98 ) from exc
101def read_json_file(raw: str) -> dict[str, Any]:
102 """Read one bounded regular JSON file without following symlinks."""
103 path = Path(raw)
104 try:
105 if path.is_symlink():
106 raise ResourceToolError("refusing to read a symbolic link")
107 flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0)
108 with os.fdopen(os.open(path, flags), "rb") as stream:
109 info = os.fstat(stream.fileno())
110 if not stat.S_ISREG(info.st_mode):
111 raise ResourceToolError("snapshot input must be a regular file")
112 if info.st_size > MAX_SNAPSHOT_BYTES:
113 raise ResourceToolError(f"snapshot input exceeds {MAX_SNAPSHOT_BYTES} bytes")
114 payload = stream.read(MAX_SNAPSHOT_BYTES + 1)
115 if len(payload) > MAX_SNAPSHOT_BYTES:
116 raise ResourceToolError(
117 f"snapshot input exceeds {MAX_SNAPSHOT_BYTES} bytes"
118 )
119 except ResourceToolError:
120 raise
121 except OSError as exc:
122 raise ResourceToolError(
123 f"could not read snapshot ({exc.__class__.__name__})"
124 ) from exc
125 try:
126 parsed = json.loads(payload)
127 except (RecursionError, UnicodeDecodeError, ValueError) as exc:
128 raise ResourceToolError("snapshot is not valid UTF-8 JSON") from exc
129 if not isinstance(parsed, dict):
130 raise ResourceToolError("snapshot JSON root must be an object")
131 return parsed
134def bounded_int(
135 value: Any,
136 *,
137 minimum: int = 0,
138 maximum: int = MAX_BYTES,
139) -> int | None:
140 """Return a bounded integer without accepting booleans or floats."""
141 if isinstance(value, bool):
142 return None
143 if isinstance(value, int) and minimum <= value <= maximum:
144 return value
145 return None
148def bounded_number(
149 value: Any,
150 *,
151 minimum: float = 0.0,
152 maximum: float = float(MAX_BYTES),
153) -> float | None:
154 """Return a finite bounded number without accepting booleans."""
155 if isinstance(value, bool) or not isinstance(value, (int, float)):
156 return None
157 try:
158 number = float(value)
159 except OverflowError:
160 return None
161 if not math.isfinite(number) or not minimum <= number <= maximum:
162 return None
163 return number
166def argparse_positive_int(
167 value: str,
168 *,
169 maximum: int = MAX_TASKS,
170) -> int:
171 """Argparse converter for bounded positive integers."""
172 try:
173 parsed = int(value, 10)
174 except ValueError as exc:
175 raise argparse.ArgumentTypeError("must be an integer") from exc
176 if not 1 <= parsed <= maximum:
177 raise argparse.ArgumentTypeError(f"must be between 1 and {maximum}")
178 return parsed
181def argparse_nonnegative_int(
182 value: str,
183 *,
184 maximum: int = MAX_BYTES,
185) -> int:
186 """Argparse converter for bounded nonnegative integers."""
187 try:
188 parsed = int(value, 10)
189 except ValueError as exc:
190 raise argparse.ArgumentTypeError("must be an integer") from exc
191 if not 0 <= parsed <= maximum:
192 raise argparse.ArgumentTypeError(f"must be between 0 and {maximum}")
193 return parsed
196def cli_error(parser: argparse.ArgumentParser, exc: Exception) -> int:
197 """Print one safe CLI error without a traceback."""
198 parser.exit(2, f"{parser.prog}: error: {exc}\n")
199 return 2