Skills › Scientific Agent Skills › geniml › _common.py
_common.py
407 lines · skills/geniml/scripts/_common.py at 5f2bc03
1#!/usr/bin/env python3
2"""Shared, dependency-free safety helpers for local Geniml skill CLIs."""
4from __future__ import annotations
6import csv
7import gzip
8import hashlib
9import io
10import json
11import math
12import os
13import re
14import stat
15import sys
16from pathlib import Path
17from typing import Any, Iterator
20HARD_MAX_FILES = 100_000
21HARD_MAX_BYTES = 8 * 1024**3
22HARD_MAX_RECORDS = 10_000_000
23HARD_MAX_LINE_BYTES = 1024 * 1024
24HARD_MAX_WORKERS = 256
25HARD_MAX_EPOCHS = 100_000
26MAX_COORDINATE = 2**63 - 1
28_SCHEMES = (
29 "http:",
30 "https:",
31 "ftp:",
32 "file:",
33 "s3:",
34 "gs:",
35 "hf:",
36 "ssh:",
37)
38_INTEGER = re.compile(r"^-?(?:0|[1-9][0-9]*)$")
39_FLOAT = re.compile(
40 r"^-?(?:(?:0|[1-9][0-9]*)\.[0-9]+|(?:0|[1-9][0-9]*)(?:[eE][+-]?[0-9]+))$"
41)
42_YAML_KEY = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$")
45class SafetyError(ValueError):
46 """Raised when a local-input safety contract is violated."""
49def bounded_int(
50 value: str,
51 *,
52 minimum: int = 0,
53 maximum: int,
54 label: str,
55) -> int:
56 """Parse an integer while enforcing an explicit hard bound."""
57 try:
58 parsed = int(value)
59 except ValueError as exc:
60 raise SafetyError(f"{label} must be an integer") from exc
61 if parsed < minimum or parsed > maximum:
62 raise SafetyError(f"{label} must be between {minimum} and {maximum}")
63 return parsed
66def int_type(*, minimum: int = 0, maximum: int, label: str):
67 """Return an argparse-compatible bounded integer parser."""
69 def parse(value: str) -> int:
70 try:
71 return bounded_int(
72 value,
73 minimum=minimum,
74 maximum=maximum,
75 label=label,
76 )
77 except SafetyError as exc:
78 raise ValueError(str(exc)) from exc
80 return parse
83def _reject_unsafe_text_path(raw: str) -> None:
84 if not raw or "\x00" in raw:
85 raise SafetyError("path must be a nonempty string without NUL bytes")
86 lowered = raw.strip().lower()
87 if lowered.startswith(_SCHEMES) or "://" in lowered:
88 raise SafetyError("URLs and URI-like paths are not allowed")
89 if raw.startswith("~"):
90 raise SafetyError("home expansion is not allowed; pass an explicit path")
91 if ".." in Path(raw).parts:
92 raise SafetyError("parent traversal ('..') is not allowed")
95def _reject_symlink_components(path: Path) -> None:
96 current = Path(path.anchor)
97 for part in path.parts[1:]:
98 current = current / part
99 try:
100 mode = current.lstat().st_mode
101 except FileNotFoundError:
102 continue
103 if stat.S_ISLNK(mode):
104 raise SafetyError("symlink path component is not allowed")
107def local_path(
108 raw: str,
109 *,
110 must_exist: bool = True,
111 kind: str = "any",
112) -> Path:
113 """Resolve a strict local path without following symlinks."""
114 _reject_unsafe_text_path(raw)
115 candidate = Path(raw)
116 if not candidate.is_absolute():
117 candidate = Path.cwd() / candidate
118 candidate = Path(os.path.abspath(candidate))
119 _reject_symlink_components(candidate)
121 try:
122 mode = candidate.lstat().st_mode
123 except FileNotFoundError:
124 if must_exist:
125 raise SafetyError("required local path does not exist")
126 parent = candidate.parent
127 try:
128 parent_mode = parent.lstat().st_mode
129 except FileNotFoundError as exc:
130 raise SafetyError("output parent does not exist") from exc
131 if not stat.S_ISDIR(parent_mode):
132 raise SafetyError("output parent is not a directory")
133 return candidate
135 if stat.S_ISLNK(mode):
136 raise SafetyError("symlinks are not allowed")
137 if kind == "file" and not stat.S_ISREG(mode):
138 raise SafetyError("expected a regular file")
139 if kind == "dir" and not stat.S_ISDIR(mode):
140 raise SafetyError("expected a directory")
141 if kind == "any" and not (stat.S_ISREG(mode) or stat.S_ISDIR(mode)):
142 raise SafetyError("only regular files and directories are allowed")
143 return candidate
146def _open_binary_nofollow(path: Path):
147 flags = os.O_RDONLY
148 if hasattr(os, "O_NOFOLLOW"):
149 flags |= os.O_NOFOLLOW
150 descriptor = os.open(path, flags)
151 mode = os.fstat(descriptor).st_mode
152 if not stat.S_ISREG(mode):
153 os.close(descriptor)
154 raise SafetyError("expected a regular file")
155 return os.fdopen(descriptor, "rb")
158def iter_text_lines(
159 path: Path,
160 *,
161 max_bytes: int,
162 max_records: int,
163 max_line_bytes: int = HARD_MAX_LINE_BYTES,
164) -> Iterator[tuple[int, str]]:
165 """Yield UTF-8 text lines with compressed and expanded bounds."""
166 if max_bytes < 1 or max_bytes > HARD_MAX_BYTES:
167 raise SafetyError("max_bytes is outside the hard safety bound")
168 if max_records < 1 or max_records > HARD_MAX_RECORDS:
169 raise SafetyError("max_records is outside the hard safety bound")
170 if max_line_bytes < 1 or max_line_bytes > HARD_MAX_LINE_BYTES:
171 raise SafetyError("max_line_bytes is outside the hard safety bound")
172 if path.stat().st_size > max_bytes:
173 raise SafetyError("compressed/input file exceeds byte limit")
175 expanded_bytes = 0
176 with _open_binary_nofollow(path) as raw_handle:
177 if path.name.lower().endswith(".gz"):
178 stream = gzip.GzipFile(fileobj=raw_handle, mode="rb")
179 else:
180 stream = raw_handle
181 try:
182 line_number = 0
183 while raw_line := stream.readline(max_line_bytes + 1):
184 line_number += 1
185 if line_number > max_records:
186 raise SafetyError("record limit exceeded")
187 if len(raw_line) > max_line_bytes:
188 raise SafetyError(f"line {line_number} exceeds line-size limit")
189 expanded_bytes += len(raw_line)
190 if expanded_bytes > max_bytes:
191 raise SafetyError("expanded text exceeds byte limit")
192 if b"\x00" in raw_line:
193 raise SafetyError(f"NUL byte found at line {line_number}")
194 try:
195 text = raw_line.decode("utf-8")
196 except UnicodeDecodeError as exc:
197 raise SafetyError(
198 f"file is not valid UTF-8 near line {line_number}"
199 ) from exc
200 yield line_number, text.rstrip("\r\n")
201 finally:
202 if stream is not raw_handle:
203 stream.close()
206def sha256_file(path: Path, *, max_bytes: int) -> tuple[str, int]:
207 """Hash a bounded regular file without following symlinks."""
208 if max_bytes < 1 or max_bytes > HARD_MAX_BYTES:
209 raise SafetyError("max_bytes is outside the hard safety bound")
210 digest = hashlib.sha256()
211 total = 0
212 with _open_binary_nofollow(path) as handle:
213 while True:
214 chunk = handle.read(1024 * 1024)
215 if not chunk:
216 break
217 total += len(chunk)
218 if total > max_bytes:
219 raise SafetyError("file exceeds hash byte limit")
220 digest.update(chunk)
221 return digest.hexdigest(), total
224def display_path(path: Path, index: int, mode: str) -> str:
225 """Render a path using the requested disclosure level."""
226 if mode == "redacted":
227 return f"file_{index:04d}"
228 if mode == "basename":
229 return path.name
230 if mode == "full":
231 return str(path)
232 raise SafetyError(f"unknown path display mode: {mode}")
235def add_path_mode_argument(parser) -> None:
236 parser.add_argument(
237 "--path-mode",
238 choices=("redacted", "basename", "full"),
239 default="redacted",
240 help="Path disclosure in output (default: redacted).",
241 )
244def print_json(payload: dict[str, Any]) -> None:
245 """Print deterministic JSON with no non-finite values."""
246 json.dump(payload, sys.stdout, indent=2, sort_keys=True, allow_nan=False)
247 sys.stdout.write("\n")
250def read_delimited_manifest(
251 path: Path,
252 *,
253 delimiter_name: str,
254 max_bytes: int,
255 max_rows: int,
256) -> tuple[list[str], list[dict[str, str]]]:
257 """Read a bounded UTF-8 CSV/TSV manifest."""
258 delimiter = "\t" if delimiter_name == "tsv" else ","
259 lines = [
260 line
261 for _, line in iter_text_lines(
262 path,
263 max_bytes=max_bytes,
264 max_records=max_rows + 1,
265 )
266 ]
267 if not lines:
268 raise SafetyError("manifest is empty")
269 reader = csv.DictReader(io.StringIO("\n".join(lines)), delimiter=delimiter)
270 if not reader.fieldnames:
271 raise SafetyError("manifest has no header")
272 raw_fieldnames = [field if field is not None else "" for field in reader.fieldnames]
273 if any(field != field.strip() for field in raw_fieldnames):
274 raise SafetyError("manifest column names must not have surrounding whitespace")
275 fieldnames = raw_fieldnames
276 if any(not field for field in fieldnames):
277 raise SafetyError("manifest contains an empty column name")
278 if len(set(fieldnames)) != len(fieldnames):
279 raise SafetyError("manifest contains duplicate column names")
280 if len(fieldnames) > 1_000:
281 raise SafetyError("manifest exceeds the 1,000-column safety bound")
283 rows: list[dict[str, str]] = []
284 for row_number, row in enumerate(reader, start=2):
285 if len(rows) >= max_rows:
286 raise SafetyError("manifest row limit exceeded")
287 if None in row:
288 raise SafetyError(f"manifest row {row_number} has extra fields")
289 normalized = {
290 key.strip(): (value.strip() if value is not None else "")
291 for key, value in row.items()
292 }
293 rows.append(normalized)
294 return fieldnames, rows
297def delimiter_from_path(path: Path, requested: str) -> str:
298 if requested != "auto":
299 return requested
300 return "csv" if path.suffix.lower() == ".csv" else "tsv"
303def load_chrom_sizes(
304 path: Path,
305 *,
306 max_bytes: int,
307 max_records: int,
308) -> tuple[dict[str, int], list[str]]:
309 """Load a strict two-column chromosome-sizes file."""
310 sizes: dict[str, int] = {}
311 order: list[str] = []
312 for line_number, line in iter_text_lines(
313 path,
314 max_bytes=max_bytes,
315 max_records=max_records,
316 ):
317 if not line or line.startswith("#"):
318 continue
319 fields = line.split("\t")
320 if len(fields) != 2:
321 raise SafetyError(
322 f"chromosome-sizes line {line_number} must have two tab-separated fields"
323 )
324 chrom, size_text = fields
325 if not chrom or any(character.isspace() for character in chrom):
326 raise SafetyError(f"invalid contig name at line {line_number}")
327 if chrom in sizes:
328 raise SafetyError(f"duplicate contig in chromosome sizes: {chrom}")
329 if not _INTEGER.fullmatch(size_text):
330 raise SafetyError(f"invalid contig size at line {line_number}")
331 size = int(size_text)
332 if size <= 0 or size > MAX_COORDINATE:
333 raise SafetyError(f"contig size out of bounds at line {line_number}")
334 sizes[chrom] = size
335 order.append(chrom)
336 if not sizes:
337 raise SafetyError("chromosome-sizes file has no records")
338 return sizes, order
341def simple_yaml_mapping(
342 path: Path,
343 *,
344 max_bytes: int,
345 max_records: int = 10_000,
346) -> dict[str, Any]:
347 """Parse a conservative top-level scalar YAML mapping without PyYAML."""
348 result: dict[str, Any] = {}
349 for line_number, line in iter_text_lines(
350 path,
351 max_bytes=max_bytes,
352 max_records=max_records,
353 ):
354 stripped = line.strip()
355 if not stripped or stripped.startswith("#") or stripped == "---":
356 continue
357 if line[:1].isspace():
358 raise SafetyError("nested YAML is not supported by the safe inspector")
359 if ":" not in line:
360 raise SafetyError(f"unsupported YAML at line {line_number}")
361 key, raw_value = line.split(":", 1)
362 key = key.strip()
363 raw_value = raw_value.strip()
364 if not _YAML_KEY.fullmatch(key):
365 raise SafetyError(f"invalid YAML key at line {line_number}")
366 if key in result:
367 raise SafetyError(f"duplicate YAML key: {key}")
368 if any(marker in raw_value for marker in ("!", "&", "*", "{", "[", "|", ">")):
369 raise SafetyError(f"complex YAML is not supported at line {line_number}")
370 if not raw_value or raw_value in {"null", "Null", "NULL", "~"}:
371 value: Any = None
372 elif raw_value in {"true", "True", "TRUE"}:
373 value = True
374 elif raw_value in {"false", "False", "FALSE"}:
375 value = False
376 elif _INTEGER.fullmatch(raw_value):
377 try:
378 value = int(raw_value)
379 except ValueError as exc:
380 raise SafetyError("YAML integer exceeds the supported size") from exc
381 elif _FLOAT.fullmatch(raw_value):
382 value = float(raw_value)
383 if not math.isfinite(value):
384 raise SafetyError("YAML numeric metadata must be finite")
385 elif (
386 len(raw_value) >= 2
387 and raw_value[0] == raw_value[-1]
388 and raw_value[0] in {"'", '"'}
389 ):
390 value = raw_value[1:-1]
391 else:
392 value = raw_value
393 result[key] = value
394 return result
397def fail_json(tool: str, exc: Exception) -> int:
398 """Emit a bounded machine-readable error."""
399 print_json(
400 {
401 "ok": False,
402 "tool": tool,
403 "error": type(exc).__name__,
404 "message": str(exc)[:500],
405 }
406 )
407 return 2