Skills › Scientific Agent Skills › gtars › _common.py
_common.py
476 lines · skills/gtars/scripts/_common.py at 5f2bc03
1#!/usr/bin/env python3
2"""Shared dependency-free safety helpers for local Gtars skill CLIs."""
4from __future__ import annotations
6import argparse
7import gzip
8import hashlib
9import json
10import os
11import re
12import stat
13import sys
14from collections import Counter
15from pathlib import Path
16from typing import Any, Iterator
19HARD_MAX_BYTES = 8 * 1024**3
20HARD_MAX_RECORDS = 10_000_000
21HARD_MAX_FILES = 100_000
22HARD_MAX_WORKERS = 256
23HARD_MAX_LINE_BYTES = 1024 * 1024
24MAX_COORDINATE = 2**32 - 1
26_SCHEMES = (
27 "http:",
28 "https:",
29 "ftp:",
30 "file:",
31 "s3:",
32 "gs:",
33 "hf:",
34 "ssh:",
35)
36_INTEGER = re.compile(r"^(?:0|[1-9][0-9]*)$")
37_VALID_STRANDS = {"+", "-", "."}
40class SafetyError(ValueError):
41 """Raised when a local-input safety contract is violated."""
44def bounded_int(
45 value: str,
46 *,
47 minimum: int = 0,
48 maximum: int,
49 label: str,
50) -> int:
51 """Parse an integer while enforcing an explicit hard bound."""
52 try:
53 parsed = int(value)
54 except ValueError as exc:
55 raise argparse.ArgumentTypeError(f"{label} must be an integer") from exc
56 if parsed < minimum or parsed > maximum:
57 raise argparse.ArgumentTypeError(
58 f"{label} must be between {minimum} and {maximum}"
59 )
60 return parsed
63def int_type(*, minimum: int = 0, maximum: int, label: str):
64 """Return an argparse-compatible bounded integer parser."""
66 def parse(value: str) -> int:
67 return bounded_int(
68 value,
69 minimum=minimum,
70 maximum=maximum,
71 label=label,
72 )
74 return parse
77def _reject_unsafe_text_path(raw: str) -> None:
78 if not raw or "\x00" in raw:
79 raise SafetyError("path must be nonempty and contain no NUL byte")
80 lowered = raw.strip().lower()
81 if lowered.startswith(_SCHEMES) or "://" in lowered:
82 raise SafetyError("URLs and URI-like paths are not allowed")
83 if raw.startswith("~"):
84 raise SafetyError("home expansion is not allowed; pass an explicit path")
85 if ".." in Path(raw).parts:
86 raise SafetyError("parent traversal ('..') is not allowed")
89def _reject_symlink_components(path: Path) -> None:
90 current = Path(path.anchor)
91 for part in path.parts[1:]:
92 current = current / part
93 try:
94 mode = current.lstat().st_mode
95 except FileNotFoundError:
96 continue
97 if stat.S_ISLNK(mode):
98 raise SafetyError("symlink path component is not allowed")
101def local_path(
102 raw: str,
103 *,
104 must_exist: bool = True,
105 kind: str = "any",
106) -> Path:
107 """Resolve a strict local path without following symlinks."""
108 _reject_unsafe_text_path(raw)
109 candidate = Path(raw)
110 if not candidate.is_absolute():
111 candidate = Path.cwd() / candidate
112 candidate = Path(os.path.abspath(candidate))
113 _reject_symlink_components(candidate)
115 try:
116 mode = candidate.lstat().st_mode
117 except FileNotFoundError:
118 if must_exist:
119 raise SafetyError("required local path does not exist")
120 parent = candidate.parent
121 try:
122 parent_mode = parent.lstat().st_mode
123 except FileNotFoundError as exc:
124 raise SafetyError("output parent does not exist") from exc
125 if not stat.S_ISDIR(parent_mode):
126 raise SafetyError("output parent is not a directory")
127 return candidate
129 if stat.S_ISLNK(mode):
130 raise SafetyError("symlinks are not allowed")
131 if kind == "file" and not stat.S_ISREG(mode):
132 raise SafetyError("expected a regular file")
133 if kind == "dir" and not stat.S_ISDIR(mode):
134 raise SafetyError("expected a directory")
135 if kind == "any" and not (stat.S_ISREG(mode) or stat.S_ISDIR(mode)):
136 raise SafetyError("only regular files and directories are allowed")
137 return candidate
140def _open_binary_nofollow(path: Path):
141 flags = os.O_RDONLY
142 if hasattr(os, "O_NOFOLLOW"):
143 flags |= os.O_NOFOLLOW
144 descriptor = os.open(path, flags)
145 mode = os.fstat(descriptor).st_mode
146 if not stat.S_ISREG(mode):
147 os.close(descriptor)
148 raise SafetyError("expected a regular file")
149 return os.fdopen(descriptor, "rb")
152def iter_text_lines(
153 path: Path,
154 *,
155 max_bytes: int,
156 max_records: int,
157 max_line_bytes: int = HARD_MAX_LINE_BYTES,
158) -> Iterator[tuple[int, str]]:
159 """Yield bounded UTF-8 lines from a local plain-text or gzip file."""
160 if not 1 <= max_bytes <= HARD_MAX_BYTES:
161 raise SafetyError("max_bytes is outside the hard safety bound")
162 if not 1 <= max_records <= HARD_MAX_RECORDS:
163 raise SafetyError("max_records is outside the hard safety bound")
164 if not 1 <= max_line_bytes <= HARD_MAX_LINE_BYTES:
165 raise SafetyError("max_line_bytes is outside the hard safety bound")
166 if path.stat().st_size > max_bytes:
167 raise SafetyError("compressed/input file exceeds byte limit")
169 expanded = 0
170 with _open_binary_nofollow(path) as raw_handle:
171 stream = (
172 gzip.GzipFile(fileobj=raw_handle, mode="rb")
173 if path.name.lower().endswith(".gz")
174 else raw_handle
175 )
176 try:
177 # Bound the read itself: iteration would allocate an entire malformed
178 # or decompressed line before the length check could reject it.
179 lines = iter(lambda: stream.readline(max_line_bytes + 1), b"")
180 for line_number, raw_line in enumerate(lines, start=1):
181 if line_number > max_records:
182 raise SafetyError("record limit exceeded")
183 if len(raw_line) > max_line_bytes:
184 raise SafetyError(f"line {line_number} exceeds line-size limit")
185 expanded += len(raw_line)
186 if expanded > max_bytes:
187 raise SafetyError("expanded text exceeds byte limit")
188 if b"\x00" in raw_line:
189 raise SafetyError(f"NUL byte found at line {line_number}")
190 try:
191 line = raw_line.decode("utf-8")
192 except UnicodeDecodeError as exc:
193 raise SafetyError(
194 f"input is not UTF-8 near line {line_number}"
195 ) from exc
196 yield line_number, line.rstrip("\r\n")
197 finally:
198 if stream is not raw_handle:
199 stream.close()
202def sha256_file(path: Path, *, max_bytes: int) -> tuple[str, int]:
203 """Hash a bounded regular file without following symlinks."""
204 if not 1 <= max_bytes <= HARD_MAX_BYTES:
205 raise SafetyError("max_bytes is outside the hard safety bound")
206 digest = hashlib.sha256()
207 total = 0
208 with _open_binary_nofollow(path) as handle:
209 while True:
210 chunk = handle.read(1024 * 1024)
211 if not chunk:
212 break
213 total += len(chunk)
214 if total > max_bytes:
215 raise SafetyError("file exceeds hash byte limit")
216 digest.update(chunk)
217 return digest.hexdigest(), total
220def prefix_collision_count(prefix: Path) -> int:
221 """Conservatively reserve a complete output prefix, including derived files."""
222 collisions = 0
223 with os.scandir(prefix.parent) as entries:
224 for index, entry in enumerate(entries, start=1):
225 if index > HARD_MAX_FILES:
226 raise SafetyError("output directory exceeds file-count inspection limit")
227 if entry.name.startswith(prefix.name):
228 collisions += 1
229 return collisions
232def load_json(path: Path, *, max_bytes: int) -> Any:
233 """Load strict bounded JSON, rejecting duplicate keys and non-finite values."""
234 text = "\n".join(
235 line
236 for _, line in iter_text_lines(
237 path,
238 max_bytes=max_bytes,
239 max_records=min(HARD_MAX_RECORDS, 1_000_000),
240 )
241 )
243 def object_pairs(pairs):
244 result = {}
245 for key, value in pairs:
246 if key in result:
247 raise SafetyError(f"duplicate JSON key: {key}")
248 result[key] = value
249 return result
251 def invalid_constant(value: str):
252 raise SafetyError(f"non-finite JSON number is not allowed: {value}")
254 try:
255 return json.loads(
256 text,
257 object_pairs_hook=object_pairs,
258 parse_constant=invalid_constant,
259 )
260 except json.JSONDecodeError as exc:
261 raise SafetyError("invalid JSON") from exc
264def load_chrom_sizes(
265 path: Path,
266 *,
267 max_bytes: int,
268 max_records: int,
269) -> tuple[dict[str, int], list[str]]:
270 """Load a strict two-column chromosome-sizes file."""
271 sizes: dict[str, int] = {}
272 order: list[str] = []
273 for line_number, line in iter_text_lines(
274 path,
275 max_bytes=max_bytes,
276 max_records=max_records,
277 ):
278 if not line or line.startswith("#"):
279 continue
280 fields = line.split("\t")
281 if len(fields) != 2:
282 raise SafetyError(
283 f"chromosome-sizes line {line_number} must have two tab fields"
284 )
285 chrom, size_text = fields
286 if (
287 not chrom
288 or any(character.isspace() or ord(character) < 32 for character in chrom)
289 or chrom in sizes
290 ):
291 raise SafetyError(f"invalid or duplicate contig at line {line_number}")
292 if not _INTEGER.fullmatch(size_text):
293 raise SafetyError(f"invalid contig size at line {line_number}")
294 size = int(size_text)
295 if not 1 <= size <= MAX_COORDINATE:
296 raise SafetyError(f"contig size out of bounds at line {line_number}")
297 sizes[chrom] = size
298 order.append(chrom)
299 if not sizes:
300 raise SafetyError("chromosome-sizes file contains no records")
301 return sizes, order
304def inspect_bed(
305 path: Path,
306 *,
307 max_bytes: int,
308 max_records: int,
309 chrom_sizes: dict[str, int] | None = None,
310 chrom_order: list[str] | None = None,
311 max_examples: int = 10,
312) -> dict[str, Any]:
313 """Inspect BED-like records without returning genomic coordinates."""
314 errors: Counter[str] = Counter()
315 warnings: Counter[str] = Counter()
316 examples: list[dict[str, int | str]] = []
317 seen: set[tuple[str, int, int]] = set()
318 max_end: dict[str, int] = {}
319 records = 0
320 skipped = 0
321 contigs: set[str] = set()
322 min_columns: int | None = None
323 max_columns = 0
324 unsorted = 0
325 duplicates = 0
326 overlap_observations = 0
327 bed6_rows = 0
328 previous_key: tuple[int | str, int, int] | None = None
329 order_map = (
330 {chrom: index for index, chrom in enumerate(chrom_order)}
331 if chrom_order is not None
332 else None
333 )
335 def issue(kind: str, code: str, line_number: int) -> None:
336 target = errors if kind == "error" else warnings
337 target[code] += 1
338 if len(examples) < max_examples:
339 examples.append({"severity": kind, "code": code, "line": line_number})
341 for line_number, line in iter_text_lines(
342 path,
343 max_bytes=max_bytes,
344 max_records=max_records,
345 ):
346 if not line:
347 skipped += 1
348 warnings["blank_line"] += 1
349 continue
350 if line.startswith(("#", "track", "browser")):
351 skipped += 1
352 continue
353 fields = line.split("\t")
354 if len(fields) < 3:
355 issue("error", "fewer_than_three_columns", line_number)
356 continue
357 records += 1
358 min_columns = len(fields) if min_columns is None else min(min_columns, len(fields))
359 max_columns = max(max_columns, len(fields))
360 chrom = fields[0]
361 if not chrom or any(
362 character.isspace() or ord(character) < 32 for character in chrom
363 ):
364 issue("error", "invalid_contig", line_number)
365 continue
366 if not _INTEGER.fullmatch(fields[1]) or not _INTEGER.fullmatch(fields[2]):
367 issue("error", "non_unsigned_decimal_coordinate", line_number)
368 continue
369 start, end = int(fields[1]), int(fields[2])
370 if start > MAX_COORDINATE or end > MAX_COORDINATE:
371 issue("error", "coordinate_exceeds_gtars_u32", line_number)
372 continue
373 if end <= start:
374 issue(
375 "error",
376 "zero_length_interval" if end == start else "end_before_start",
377 line_number,
378 )
379 continue
380 if chrom_sizes is not None:
381 if chrom not in chrom_sizes:
382 issue("error", "unknown_contig", line_number)
383 continue
384 if end > chrom_sizes[chrom]:
385 issue("error", "end_beyond_contig", line_number)
386 continue
387 if len(fields) >= 6:
388 bed6_rows += 1
389 if fields[5] not in _VALID_STRANDS:
390 issue("error", "invalid_bed6_strand", line_number)
392 contigs.add(chrom)
393 key = (chrom, start, end)
394 if key in seen:
395 duplicates += 1
396 else:
397 seen.add(key)
398 prior_end = max_end.get(chrom)
399 if prior_end is not None and start < prior_end:
400 overlap_observations += 1
401 max_end[chrom] = max(end, prior_end or end)
403 sort_key: tuple[int | str, int, int]
404 if order_map is None:
405 sort_key = (chrom, start, end)
406 else:
407 sort_key = (order_map[chrom], start, end)
408 if previous_key is not None and sort_key < previous_key:
409 unsorted += 1
410 previous_key = sort_key
412 if duplicates:
413 warnings["duplicate_intervals"] = duplicates
414 if overlap_observations:
415 warnings["overlapping_intervals"] = overlap_observations
416 if unsorted:
417 warnings["out_of_order_records"] = unsorted
418 if chrom_sizes is None:
419 warnings["bounds_not_checked_without_chrom_sizes"] += 1
421 digest, size = sha256_file(path, max_bytes=max_bytes)
422 return {
423 "sha256": digest,
424 "size_bytes": size,
425 "records": records,
426 "skipped_records": skipped,
427 "contig_count": len(contigs),
428 "minimum_columns": min_columns,
429 "maximum_columns": max_columns if records else None,
430 "bed6_or_wider_rows": bed6_rows,
431 "duplicate_intervals": duplicates,
432 "overlap_observations": overlap_observations,
433 "out_of_order_records": unsorted,
434 "errors": dict(sorted(errors.items())),
435 "warnings": dict(sorted(warnings.items())),
436 "issue_examples": examples,
437 }
440def add_path_mode_argument(parser: argparse.ArgumentParser) -> None:
441 parser.add_argument(
442 "--path-mode",
443 choices=("redacted", "basename", "full"),
444 default="redacted",
445 help="Path disclosure in output (default: redacted).",
446 )
449def display_path(path: Path, index: int, mode: str) -> str:
450 """Render a path using the requested disclosure level."""
451 if mode == "redacted":
452 return f"file_{index:04d}"
453 if mode == "basename":
454 return path.name
455 if mode == "full":
456 return str(path)
457 raise SafetyError(f"unknown path mode: {mode}")
460def print_json(payload: dict[str, Any]) -> None:
461 """Print deterministic JSON with no non-finite values."""
462 json.dump(payload, sys.stdout, indent=2, sort_keys=True, allow_nan=False)
463 sys.stdout.write("\n")
466def fail_json(tool: str, exc: Exception) -> int:
467 """Emit a bounded machine-readable error."""
468 print_json(
469 {
470 "ok": False,
471 "tool": tool,
472 "error": type(exc).__name__,
473 "message": str(exc)[:500],
474 }
475 )
476 return 2