penrun 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
penrun-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,7 @@
1
+ Metadata-Version: 2.3
2
+ Name: penrun
3
+ Version: 0.1.0
4
+ Summary: Run any command inside a well-defined, timestamped artifacts directory.
5
+ Author: AISEC Pentesting Team
6
+ License: Apache-2.0
7
+ Requires-Python: >=3.14
@@ -0,0 +1,62 @@
1
+ # SPDX-FileCopyrightText: AISEC Pentesting Team
2
+ #
3
+ # SPDX-License-Identifier: CC0-1.0
4
+
5
+ [build-system]
6
+ requires = ["uv_build>=0.9.11,<0.12.0"]
7
+ build-backend = "uv_build"
8
+
9
+ [project]
10
+ name = "penrun"
11
+ version = "0.1.0"
12
+ description = "Run any command inside a well-defined, timestamped artifacts directory."
13
+ requires-python = ">=3.14"
14
+ license = { text = "Apache-2.0" }
15
+ authors = [{ name = "AISEC Pentesting Team" }]
16
+
17
+ [project.scripts]
18
+ penrun = "penrun.cli:main"
19
+
20
+ [dependency-groups]
21
+ dev = [
22
+ "ruff",
23
+ "mypy",
24
+ "ty",
25
+ "reuse>=6.2.0",
26
+ "rust-just>=1.38.0",
27
+ ]
28
+
29
+ [tool.ruff]
30
+ target-version = "py314"
31
+ line-length = 110
32
+
33
+ [tool.ruff.lint]
34
+ select = [
35
+ "B", # flake8-bugbear
36
+ "A", # flake8-builtins
37
+ "ASYNC",# flake8-async
38
+ "C4", # flake8-comprehensions
39
+ "E", # pycodestlye
40
+ "F", # pyflakes
41
+ "FA", # flake8-future-annotations
42
+ "FLY", # flynt
43
+ "FURB", # refurb
44
+ "I", # isort
45
+ "IC", # flake8-import-conventions
46
+ "LOG", # flake8-logging
47
+ "PGH", # pygrep-hooks
48
+ "PIE", # flake8-pie
49
+ "PL", # pylint
50
+ "PTH", # flake8-use-pathlib
51
+ "Q", # flake8-quotes
52
+ "RSE", # flake8-raise
53
+ "RUF100", # unused-noqa
54
+ "TD", # flake8-todos
55
+ "TID", # flake8-tidy-imports
56
+ "UP", # pyupgrade
57
+ ]
58
+
59
+ [tool.mypy]
60
+ strict = true
61
+ native_parser = true
62
+ num_workers = 4
@@ -0,0 +1,10 @@
1
+ # SPDX-FileCopyrightText: AISEC Pentesting Team
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from api import run, run_batched
6
+
7
+ __all__ = [
8
+ "run",
9
+ "run_batched",
10
+ ]
@@ -0,0 +1,10 @@
1
+ # SPDX-FileCopyrightText: AISEC Pentesting Team
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import sys
6
+
7
+ from penrun.cli import main
8
+
9
+ if __name__ == "__main__":
10
+ sys.exit(main())
@@ -0,0 +1,84 @@
1
+ # SPDX-FileCopyrightText: AISEC Pentesting Team
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from contextlib import nullcontext
6
+ from dataclasses import dataclass
7
+ from pathlib import Path
8
+
9
+ from penrun.batch import BatchOptions, BatchRunner, InputItem
10
+ from penrun.config import Config
11
+ from penrun.locking import FileLock
12
+ from penrun.runner import Runner, RunRequest
13
+
14
+ __all__ = ["BatchOptions", "InputItem", "InvocationData", "RunOptions", "run", "run_batched"]
15
+
16
+
17
+ @dataclass(frozen=True, slots=True)
18
+ class RunOptions:
19
+ config: Config | None = None
20
+ tag: str = ""
21
+ artifactsdir_override: Path | None = None
22
+ base_override: Path | None = None
23
+ chdir: bool = False
24
+ include_env: bool = False
25
+
26
+
27
+ @dataclass(frozen=True, slots=True)
28
+ class InvocationData:
29
+ # The batch's shared executable and args. Never executed directly -
30
+ # each job's own InputItem.values is already the complete, resolved
31
+ # command - but arg0/arg0_name give every job in the batch the same
32
+ # {scriptname} artifacts_template placeholder, and argv is available
33
+ # for callers that want to log/display the batch's base invocation.
34
+ # Turning raw stdin lines/a -T template into this shape is a CLI-only
35
+ # concern - see cli.py.
36
+ arg0: str
37
+ arg0_name: str | None
38
+ argv: list[str]
39
+ inputs: list[InputItem]
40
+
41
+
42
+ def _lock(config: Config) -> nullcontext[None] | FileLock:
43
+ return FileLock(Path(config.lock)) if config.lock else nullcontext()
44
+
45
+
46
+ def run(cmd: list[str], options: RunOptions | None = None, *, mirror_to_stdout: bool = True) -> int:
47
+ options = options or RunOptions()
48
+ config = options.config or Config()
49
+ request = RunRequest(
50
+ cmd=cmd,
51
+ tag=options.tag,
52
+ artifactsdir_override=options.artifactsdir_override,
53
+ base_override=options.base_override,
54
+ chdir=options.chdir,
55
+ include_env=options.include_env,
56
+ mirror_to_stdout=mirror_to_stdout,
57
+ )
58
+ with _lock(config):
59
+ return Runner(config).run(request)
60
+
61
+
62
+ def run_batched(
63
+ data: InvocationData,
64
+ options: RunOptions | None = None,
65
+ *,
66
+ named_jobs: bool = False,
67
+ batch: BatchOptions | None = None,
68
+ ) -> int:
69
+ options = options or RunOptions()
70
+ batch = batch or BatchOptions()
71
+ config = options.config or Config()
72
+ scriptname = data.arg0_name if data.arg0_name is not None else Path(data.arg0).stem
73
+ base_request = RunRequest(
74
+ cmd=[],
75
+ tag=options.tag,
76
+ artifactsdir_override=options.artifactsdir_override,
77
+ base_override=options.base_override,
78
+ chdir=options.chdir,
79
+ include_env=options.include_env,
80
+ scriptname=scriptname,
81
+ named_jobs=named_jobs,
82
+ )
83
+ with _lock(config):
84
+ return BatchRunner(Runner(config), batch, base_request).run(data.inputs)
@@ -0,0 +1,93 @@
1
+ # SPDX-FileCopyrightText: AISEC Pentesting Team
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import json
6
+ import os
7
+ from dataclasses import dataclass
8
+ from datetime import datetime
9
+ from pathlib import Path
10
+
11
+ from penrun.exceptions import ConfigError
12
+
13
+
14
+ @dataclass(frozen=True, slots=True)
15
+ class ArtifactsPlaceholders:
16
+ base: str
17
+ scriptname: str
18
+ tag: str
19
+ fragment: str
20
+ jobname: str
21
+ timestamp: str
22
+
23
+ @property
24
+ def tagsuffix(self) -> str:
25
+ return f"-{self.tag}" if self.tag else ""
26
+
27
+ def as_dict(self) -> dict[str, str]:
28
+ return {
29
+ "base": self.base,
30
+ "scriptname": self.scriptname,
31
+ "tag": self.tag,
32
+ "tagsuffix": self.tagsuffix,
33
+ "fragment": self.fragment,
34
+ "jobname": self.jobname,
35
+ "timestamp": self.timestamp,
36
+ }
37
+
38
+
39
+ def render_artifactsdir(template: str, placeholders: ArtifactsPlaceholders) -> Path:
40
+ try:
41
+ rendered = template.format(**placeholders.as_dict())
42
+ except KeyError as e:
43
+ raise ConfigError(f"invalid artifacts_template placeholder: {e}") from e
44
+ return Path(rendered)
45
+
46
+
47
+ def timestamp_now() -> str:
48
+ return datetime.now().strftime("%Y%m%d-%H%M%S.%f")
49
+
50
+
51
+ def now_iso() -> str:
52
+ return datetime.now().astimezone().isoformat(timespec="seconds")
53
+
54
+
55
+ def update_latest_symlink(link_path: Path, target: Path) -> None:
56
+ link_path.parent.mkdir(parents=True, exist_ok=True)
57
+ if link_path.is_symlink() or link_path.exists():
58
+ link_path.unlink()
59
+ link_path.symlink_to(os.path.relpath(target, link_path.parent))
60
+
61
+
62
+ class ArtifactsDir:
63
+ def __init__(self, path: Path, latest_dir: Path) -> None:
64
+ self.path = path
65
+ self._latest_dir = latest_dir
66
+
67
+ def create(self) -> None:
68
+ self.path.mkdir(parents=True, exist_ok=True)
69
+
70
+ def output_path(self) -> Path:
71
+ return self.path.joinpath("OUTPUT.zst")
72
+
73
+ def update_latest_symlink(self) -> None:
74
+ update_latest_symlink(self._latest_dir.joinpath("LATEST"), self.path)
75
+
76
+ def write_meta(
77
+ self,
78
+ *,
79
+ command: list[str],
80
+ start_time: str,
81
+ end_time: str,
82
+ exit_code: int,
83
+ env: dict[str, str] | None = None,
84
+ ) -> None:
85
+ meta = {
86
+ "command": command,
87
+ "start_time": start_time,
88
+ "end_time": end_time,
89
+ "exit_code": exit_code,
90
+ }
91
+ if env is not None:
92
+ meta["env"] = env
93
+ self.path.joinpath("META.json").write_text(json.dumps(meta) + "\n")
@@ -0,0 +1,160 @@
1
+ # SPDX-FileCopyrightText: AISEC Pentesting Team
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import shlex
6
+ import threading
7
+ import time
8
+ from collections.abc import Iterable
9
+ from concurrent.futures import ThreadPoolExecutor
10
+ from dataclasses import dataclass, replace
11
+
12
+ from penrun.artifacts import timestamp_now
13
+ from penrun.constants import EXIT_CODE_FAILED_FIRST, EXIT_CODE_JOB_FAILED
14
+ from penrun.output import log
15
+ from penrun.runner import Runner, RunRequest
16
+
17
+
18
+ def _ratio(count: int, total: int) -> str:
19
+ pct = count / total * 100 if total else 0.0
20
+ return f"{count}/{total} ({pct:.0f}%)"
21
+
22
+
23
+ def _format_duration(seconds: float) -> str:
24
+ seconds = int(seconds)
25
+ hours, seconds = divmod(seconds, 3600)
26
+ minutes, seconds = divmod(seconds, 60)
27
+ if hours:
28
+ return f"{hours}:{minutes:02d}:{seconds:02d}"
29
+ return f"{minutes}:{seconds:02d}"
30
+
31
+
32
+ @dataclass(frozen=True, slots=True)
33
+ class InputItem:
34
+ # A single batch job's fully-resolved command (values[0] is the
35
+ # executable) plus the artifacts directory name it should run under.
36
+ # Building these from raw input lines/templates is a CLI-only concern -
37
+ # see cli.py.
38
+ values: list[str]
39
+ dirname: str
40
+
41
+
42
+ @dataclass(frozen=True, slots=True)
43
+ class BatchOptions:
44
+ jobs: int = 1
45
+ sleep: float = 0.0
46
+ until: bool = False
47
+
48
+
49
+ class BatchRunner:
50
+ def __init__(self, runner: Runner, options: BatchOptions, base_request: RunRequest) -> None:
51
+ self._runner = runner
52
+ self._options = options
53
+ self._base_request = base_request
54
+
55
+ def run(self, items: Iterable[InputItem]) -> int:
56
+ # All jobs of one batch share a single timestamp, so they land
57
+ # together under one DEFAULT_BATCH_TEMPLATE run-{timestamp} dir
58
+ # instead of each picking their own.
59
+ items = list(items)
60
+ self._base_request = replace(self._base_request, batch=True, timestamp=timestamp_now())
61
+ self._update_latest(items)
62
+
63
+ if self._options.jobs > 1:
64
+ return self._run_parallel(items)
65
+ return self._run_sequential(items)
66
+
67
+ def _update_latest(self, items: list[InputItem]) -> None:
68
+ # Point a single top-level LATEST at the run-{timestamp} directory
69
+ # shared by every job of this batch, rather than each job trying
70
+ # to maintain its own (see Runner.update_batch_latest).
71
+ if items:
72
+ self._runner.update_batch_latest(self._build_request(items[0]))
73
+
74
+ def _build_request(self, item: InputItem, *, mirror_to_stdout: bool = True) -> RunRequest:
75
+ return replace(
76
+ self._base_request,
77
+ cmd=item.values,
78
+ jobname=item.dirname,
79
+ mirror_to_stdout=mirror_to_stdout,
80
+ )
81
+
82
+ def _run_sequential(self, items: Iterable[InputItem]) -> int:
83
+ result = 0
84
+ for item in items:
85
+ log(f"Starting command '{shlex.join(item.values)}'.")
86
+ ret = self._runner.run(self._build_request(item))
87
+ if ret != 0:
88
+ if self._options.until:
89
+ return EXIT_CODE_FAILED_FIRST
90
+ result = EXIT_CODE_JOB_FAILED
91
+ time.sleep(self._options.sleep)
92
+ return result
93
+
94
+ def _run_parallel(self, items: list[InputItem]) -> int:
95
+ # Live terminal output is disabled for job control: interleaved
96
+ # output from parallel jobs is unreadable anyway - only the
97
+ # compressed OUTPUT file is kept per job.
98
+ total = len(items)
99
+ start = time.monotonic()
100
+ stop = threading.Event()
101
+ lock = threading.Lock()
102
+ active = 0
103
+ finished = 0
104
+ executed = 0
105
+ succeeded = 0
106
+
107
+ def progress(prefix: str) -> None:
108
+ remaining = total - active - finished
109
+ elapsed = time.monotonic() - start
110
+ eta = _format_duration(elapsed / finished * remaining) if finished else "unknown"
111
+ success_rate = _ratio(succeeded, executed) if executed else "n/a"
112
+ log(
113
+ f"{prefix} Active: {_ratio(active, total)}, "
114
+ f"finished: {_ratio(finished, total)}, "
115
+ f"remaining: {_ratio(remaining, total)}. "
116
+ f"Runtime: {_format_duration(elapsed)}, ETA: {eta}. "
117
+ f"Success rate: {success_rate}."
118
+ )
119
+
120
+ def job(item: InputItem) -> int | None:
121
+ nonlocal active, finished, executed, succeeded
122
+ if stop.is_set():
123
+ with lock:
124
+ finished += 1
125
+ return None
126
+
127
+ with lock:
128
+ active += 1
129
+ progress(f"Scheduling job '{shlex.join(item.values)}'.")
130
+
131
+ ret = self._runner.run(self._build_request(item, mirror_to_stdout=False))
132
+
133
+ with lock:
134
+ active -= 1
135
+ finished += 1
136
+ executed += 1
137
+ if ret == 0:
138
+ succeeded += 1
139
+ progress("Job terminated.")
140
+ return ret
141
+
142
+ failed = False
143
+ with ThreadPoolExecutor(max_workers=self._options.jobs) as executor:
144
+ futures = []
145
+ for item in items:
146
+ if stop.is_set():
147
+ break
148
+ futures.append(executor.submit(job, item))
149
+ time.sleep(self._options.sleep)
150
+
151
+ for future in futures:
152
+ ret = future.result()
153
+ if ret is not None and ret != 0:
154
+ failed = True
155
+ if self._options.until:
156
+ stop.set()
157
+
158
+ if failed:
159
+ return EXIT_CODE_FAILED_FIRST if self._options.until else EXIT_CODE_JOB_FAILED
160
+ return 0
@@ -0,0 +1,242 @@
1
+ # SPDX-FileCopyrightText: AISEC Pentesting Team
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import argparse
6
+ import os
7
+ import shlex
8
+ import signal
9
+ import sys
10
+ from dataclasses import dataclass
11
+ from pathlib import Path
12
+
13
+ from penrun import api
14
+ from penrun import config as config_mod
15
+ from penrun.api import BatchOptions, InputItem, InvocationData, RunOptions
16
+ from penrun.config import Config
17
+ from penrun.exceptions import PenrunError, UsageError
18
+ from penrun.output import log
19
+
20
+ _SHORT_FLAG_LEN = len("-x")
21
+
22
+
23
+ @dataclass(frozen=True, slots=True)
24
+ class _OptionSpec:
25
+ short: str
26
+ long: str
27
+ dest: str
28
+ takes_value: bool
29
+ help: str
30
+ type: type = str
31
+ default: object = None
32
+
33
+
34
+ OPTIONS: tuple[_OptionSpec, ...] = (
35
+ _OptionSpec(
36
+ "-S", "--sleep", "sleep", True, "Sleep this time between jobs (batch mode)", type=float, default=0.0
37
+ ),
38
+ _OptionSpec(
39
+ "-T", "--template", "template", True, "A command template; {1} is the input line, may repeat"
40
+ ),
41
+ _OptionSpec("-b", "--batch", "batch", False, "Batched mode: read commands from stdin", default=False),
42
+ _OptionSpec("-B", "--base-dir", "base", True, "Use BASE as the penrun-artifacts base dir"),
43
+ _OptionSpec("-c", "--config", "config", True, "Use this config file exclusively"),
44
+ _OptionSpec("-d", "--artifacts-dir", "dir", True, "Use artifactsdir DIR"),
45
+ _OptionSpec(
46
+ "-C",
47
+ "--chdir",
48
+ "chdir",
49
+ False,
50
+ "Run COMMAND with its cwd set to the artifacts dir",
51
+ default=False,
52
+ ),
53
+ _OptionSpec(
54
+ "-j",
55
+ "--jobs",
56
+ "jobs",
57
+ True,
58
+ "Process this many jobs at a time (batch mode)",
59
+ type=int,
60
+ default=1,
61
+ ),
62
+ _OptionSpec(
63
+ "-N",
64
+ "--name-by-arg",
65
+ "name_by_fragment",
66
+ False,
67
+ "Batch mode: name each job's directory after its input line/argument instead of the program name",
68
+ default=False,
69
+ ),
70
+ _OptionSpec("-t", "--tag", "tag", True, "Add a tag to this run", default=""),
71
+ _OptionSpec("-u", "--until", "until", False, "In batch mode: stop after the first error", default=False),
72
+ )
73
+
74
+
75
+ # Splits argv into penrun's own option tokens and the trailing COMMAND,
76
+ # stopping at the first token that isn't (part of) one of our own options
77
+ # - like POSIX getopt(3)/bash's `getopts`, so `penrun -t foo ls -la` passes
78
+ # `-la` through to `ls` untouched. argparse alone can't express this (it
79
+ # permutes options and positionals freely), so this pre-scan hands it only
80
+ # the option tokens.
81
+ def _split_argv(argv: list[str]) -> tuple[list[str], list[str]]:
82
+ value_flags = {opt.short for opt in OPTIONS if opt.takes_value} | {
83
+ opt.long for opt in OPTIONS if opt.takes_value
84
+ }
85
+ bool_flags = (
86
+ {opt.short for opt in OPTIONS if not opt.takes_value}
87
+ | {opt.long for opt in OPTIONS if not opt.takes_value}
88
+ | {"-h", "--help"}
89
+ )
90
+
91
+ i = 0
92
+ n = len(argv)
93
+ while i < n:
94
+ tok = argv[i]
95
+
96
+ if tok == "--":
97
+ return argv[:i], argv[i + 1 :]
98
+ if tok in bool_flags:
99
+ i += 1
100
+ continue
101
+ if tok in value_flags:
102
+ i += 2 if i + 1 < n else 1
103
+ continue
104
+ if tok.startswith("--") and tok.partition("=")[0] in value_flags:
105
+ i += 1 # value is attached, e.g. --tag=foo
106
+ continue
107
+ if len(tok) > _SHORT_FLAG_LEN and tok[0] == "-" and tok[1] != "-" and tok[:2] in value_flags:
108
+ i += 1 # value is attached, e.g. -tfoo
109
+ continue
110
+
111
+ break
112
+
113
+ return argv[:i], argv[i:]
114
+
115
+
116
+ def _build_parser() -> argparse.ArgumentParser:
117
+ parser = argparse.ArgumentParser(
118
+ description="Run COMMAND inside a well-defined, timestamped artifacts directory.",
119
+ )
120
+ for opt in OPTIONS:
121
+ if opt.takes_value:
122
+ parser.add_argument(
123
+ opt.short,
124
+ opt.long,
125
+ dest=opt.dest,
126
+ type=opt.type,
127
+ default=opt.default,
128
+ help=opt.help,
129
+ )
130
+ else:
131
+ parser.add_argument(
132
+ opt.short,
133
+ opt.long,
134
+ dest=opt.dest,
135
+ action="store_true",
136
+ default=opt.default,
137
+ help=opt.help,
138
+ )
139
+ return parser
140
+
141
+
142
+ def _resolve_config(explicit: str | None) -> Config:
143
+ if explicit:
144
+ if explicit != "/dev/null" and not Path(explicit).is_file():
145
+ raise PenrunError("config does not exist")
146
+ for key, value in config_mod.find_config_paths().items():
147
+ os.environ[key] = str(value)
148
+ if explicit == "/dev/null":
149
+ return Config()
150
+ return config_mod.load_config(Path(explicit))
151
+
152
+ paths = config_mod.find_config_paths()
153
+ for key, value in paths.items():
154
+ os.environ[key] = str(value)
155
+ chosen = config_mod.select_config_path(paths)
156
+ return config_mod.load_config(chosen) if chosen else Config()
157
+
158
+
159
+ def _sanitize_dirname(name: str) -> str:
160
+ name = name.strip().replace(os.sep, "_")
161
+ if os.altsep:
162
+ name = name.replace(os.altsep, "_")
163
+ return name if name and name not in {".", ".."} else "job"
164
+
165
+
166
+ # The classic -T/{1} template feature: turns one raw stdin line into a full
167
+ # command, either by substituting it into the template (which may repeat
168
+ # {1} anywhere) or, without a template, by treating the line as a command
169
+ # of its own.
170
+ def _build_cmd(template: str | None, fragment: str) -> list[str] | None:
171
+ cmd_str = template.replace("{1}", fragment) if template else fragment
172
+ try:
173
+ cmd = shlex.split(cmd_str)
174
+ except ValueError as e:
175
+ log(f"error: invalid command '{cmd_str}': {e}")
176
+ return None
177
+ return cmd or None
178
+
179
+
180
+ def _read_invocation_data(ns: argparse.Namespace) -> InvocationData:
181
+ fragments = [line for line in sys.stdin.read().splitlines() if line.strip()]
182
+
183
+ inputs: list[InputItem] = []
184
+ for fragment in fragments:
185
+ cmd = _build_cmd(ns.template, fragment)
186
+ if cmd is None:
187
+ continue
188
+ dirname = _sanitize_dirname(fragment) if ns.name_by_fragment else Path(cmd[0]).stem
189
+ inputs.append(InputItem(values=cmd, dirname=dirname))
190
+
191
+ if ns.template:
192
+ arg0, *argv = shlex.split(ns.template)
193
+ elif inputs:
194
+ arg0, *argv = inputs[0].values
195
+ else:
196
+ arg0, argv = "", []
197
+
198
+ return InvocationData(arg0=arg0, arg0_name=None, argv=argv, inputs=inputs)
199
+
200
+
201
+ def _run(ns: argparse.Namespace, command: list[str]) -> int:
202
+ config = _resolve_config(ns.config)
203
+ options = RunOptions(
204
+ config=config,
205
+ tag=ns.tag,
206
+ artifactsdir_override=Path(ns.dir) if ns.dir else None,
207
+ base_override=Path(ns.base) if ns.base else None,
208
+ chdir=ns.chdir,
209
+ )
210
+
211
+ if ns.batch:
212
+ data = _read_invocation_data(ns)
213
+ os.environ["PENRUN_BATCH_MODE"] = "1"
214
+ batch = BatchOptions(jobs=ns.jobs, sleep=ns.sleep, until=ns.until)
215
+ return api.run_batched(data, options, named_jobs=ns.name_by_fragment, batch=batch)
216
+
217
+ if not command:
218
+ raise UsageError("no command specified")
219
+
220
+ return api.run(command, options)
221
+
222
+
223
+ def main(argv: list[str] | None = None) -> int:
224
+ argv = sys.argv[1:] if argv is None else argv
225
+ own_args, command = _split_argv(argv)
226
+
227
+ parser = _build_parser()
228
+ try:
229
+ namespace = parser.parse_args(own_args)
230
+ except SystemExit as e:
231
+ # argparse handles -h/--help and its own usage errors by printing
232
+ # and calling sys.exit() directly; translate that into a normal
233
+ # return so main() never exits the process itself.
234
+ return int(e.code) if isinstance(e.code, int) else 0
235
+
236
+ try:
237
+ return _run(namespace, command)
238
+ except PenrunError as e:
239
+ log(f"error: {e}")
240
+ return e.exit_code
241
+ except KeyboardInterrupt:
242
+ return 128 + signal.SIGINT
@@ -0,0 +1,109 @@
1
+ # SPDX-FileCopyrightText: AISEC Pentesting Team
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import subprocess
6
+ import tomllib
7
+ import types
8
+ import typing
9
+ from dataclasses import dataclass, fields
10
+ from pathlib import Path
11
+ from typing import Any
12
+
13
+ from penrun.exceptions import ConfigError
14
+
15
+ CONFIG_FILENAME = ".penrun.toml"
16
+ USER_CONFIG_PATH = Path("~/.config/penrun/config.toml").expanduser()
17
+
18
+ ENV_PWD_CONF = "PENRUN_PWD_CONF"
19
+ ENV_GIT_ROOT_CONF = "PENRUN_GIT_ROOT_CONF"
20
+ ENV_USER_CONF = "PENRUN_USER_CONF"
21
+
22
+ # Precedence order: first match wins, and stops the search.
23
+ _SEARCH_ORDER = (ENV_PWD_CONF, ENV_GIT_ROOT_CONF, ENV_USER_CONF)
24
+
25
+
26
+ @dataclass(frozen=True, slots=True)
27
+ class Config:
28
+ artifacts_base: str | None = None
29
+ # None means: use runner.DEFAULT_TEMPLATE, or runner.DEFAULT_BATCH_TEMPLATE
30
+ # in batch mode - see there.
31
+ artifacts_template: str | None = None
32
+ lock: str | None = None
33
+
34
+
35
+ def find_config_paths() -> dict[str, Path]:
36
+ paths: dict[str, Path] = {}
37
+
38
+ pwd_conf = Path.cwd().joinpath(CONFIG_FILENAME)
39
+ if pwd_conf.is_file():
40
+ paths[ENV_PWD_CONF] = pwd_conf
41
+
42
+ git_root_conf = _git_root_config()
43
+ if git_root_conf is not None:
44
+ paths[ENV_GIT_ROOT_CONF] = git_root_conf
45
+
46
+ if USER_CONFIG_PATH.is_file():
47
+ paths[ENV_USER_CONF] = USER_CONFIG_PATH
48
+
49
+ return paths
50
+
51
+
52
+ def select_config_path(paths: dict[str, Path]) -> Path | None:
53
+ for key in _SEARCH_ORDER:
54
+ if key in paths:
55
+ return paths[key]
56
+ return None
57
+
58
+
59
+ def load_config(path: Path) -> Config:
60
+ with path.open("rb") as f:
61
+ try:
62
+ data = tomllib.load(f)
63
+ except tomllib.TOMLDecodeError as e:
64
+ raise ConfigError(f"invalid config {path}: {e}") from e
65
+
66
+ return _build_config(path, data)
67
+
68
+
69
+ def _git_root_config() -> Path | None:
70
+ try:
71
+ result = subprocess.run(
72
+ ["git", "rev-parse", "--show-toplevel"],
73
+ capture_output=True,
74
+ text=True,
75
+ check=True,
76
+ )
77
+ except subprocess.CalledProcessError, FileNotFoundError:
78
+ return None
79
+
80
+ conf = Path(result.stdout.strip()).joinpath(CONFIG_FILENAME)
81
+ return conf if conf.is_file() else None
82
+
83
+
84
+ def _build_config(path: Path, data: dict[str, Any]) -> Config:
85
+ hints = typing.get_type_hints(Config)
86
+ allowed = {f.name for f in fields(Config)}
87
+
88
+ for key, value in data.items():
89
+ if key not in allowed:
90
+ raise ConfigError(f"invalid config {path}: unknown key '{key}'")
91
+ if not _matches_type(value, hints[key]):
92
+ raise ConfigError(f"invalid config {path}: '{key}' has the wrong type")
93
+
94
+ return Config(**data)
95
+
96
+
97
+ def _matches_type(value: Any, tp: Any) -> bool:
98
+ origin = typing.get_origin(tp)
99
+
100
+ # TOML has no null literal, so a present key is always checked against
101
+ # the non-None member of a `X | None` union.
102
+ if origin is types.UnionType:
103
+ return any(_matches_type(value, arg) for arg in typing.get_args(tp) if arg is not type(None))
104
+ if origin is list:
105
+ (elem_type,) = typing.get_args(tp) or (Any,)
106
+ return isinstance(value, list) and all(_matches_type(v, elem_type) for v in value)
107
+ if tp is str:
108
+ return isinstance(value, str)
109
+ return False
@@ -0,0 +1,12 @@
1
+ # SPDX-FileCopyrightText: AISEC Pentesting Team
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ # EX_USAGE and EX_SOFTWARE follow the BSD sysexits(3) convention:
6
+ # https://www.freebsd.org/cgi/man.cgi?query=sysexits
7
+
8
+ EXIT_CODE_JOB_FAILED = 1
9
+ EXIT_CODE_FAILED_FIRST = 2
10
+
11
+ EXIT_CODE_EX_USAGE = 64
12
+ EXIT_CODE_EX_SOFTWARE = 70
@@ -0,0 +1,26 @@
1
+ # SPDX-FileCopyrightText: AISEC Pentesting Team
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from penrun.constants import EXIT_CODE_EX_SOFTWARE, EXIT_CODE_EX_USAGE
6
+
7
+
8
+ class PenrunError(Exception):
9
+ def __init__(self, message: str, exit_code: int = EXIT_CODE_EX_SOFTWARE) -> None:
10
+ super().__init__(message)
11
+ self.exit_code = exit_code
12
+
13
+
14
+ class UsageError(PenrunError):
15
+ def __init__(self, message: str) -> None:
16
+ super().__init__(message, EXIT_CODE_EX_USAGE)
17
+
18
+
19
+ class ConfigError(PenrunError):
20
+ pass
21
+
22
+
23
+ class CommandNotFoundError(PenrunError):
24
+ def __init__(self, command: str) -> None:
25
+ super().__init__(f"command not found: {command}")
26
+ self.command = command
@@ -0,0 +1,41 @@
1
+ # SPDX-FileCopyrightText: AISEC Pentesting Team
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import fcntl
6
+ import os
7
+ from pathlib import Path
8
+ from types import TracebackType
9
+
10
+ from penrun.output import log
11
+
12
+
13
+ class FileLock:
14
+ def __init__(self, path: Path) -> None:
15
+ self._path = path
16
+ self._fd: int | None = None
17
+
18
+ def __enter__(self) -> FileLock:
19
+ self._path.parent.mkdir(parents=True, exist_ok=True)
20
+ self._path.touch(exist_ok=True)
21
+ self._fd = os.open(self._path, os.O_RDONLY)
22
+
23
+ try:
24
+ fcntl.flock(self._fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
25
+ except BlockingIOError:
26
+ log(f"Could not aquire lock {self._path}")
27
+ log("Waiting for lock…")
28
+ fcntl.flock(self._fd, fcntl.LOCK_EX)
29
+ log("… aquired lock.")
30
+
31
+ return self
32
+
33
+ def __exit__(
34
+ self,
35
+ exc_type: type[BaseException] | None,
36
+ exc: BaseException | None,
37
+ tb: TracebackType | None,
38
+ ) -> None:
39
+ if self._fd is not None:
40
+ os.close(self._fd)
41
+ self._fd = None
@@ -0,0 +1,9 @@
1
+ # SPDX-FileCopyrightText: AISEC Pentesting Team
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import sys
6
+
7
+
8
+ def log(*values: object) -> None:
9
+ print(*values, file=sys.stderr)
@@ -0,0 +1,53 @@
1
+ # SPDX-FileCopyrightText: AISEC Pentesting Team
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import io
6
+ import subprocess
7
+ import sys
8
+ from compression import zstd
9
+ from dataclasses import dataclass
10
+ from pathlib import Path
11
+
12
+ from penrun.exceptions import CommandNotFoundError
13
+
14
+
15
+ @dataclass(frozen=True, slots=True)
16
+ class PipelineSinks:
17
+ raw_out_path: Path
18
+ mirror_to_stdout: bool
19
+
20
+
21
+ def run_pipeline(cmd: list[str], env: dict[str, str], sinks: PipelineSinks, cwd: Path | None = None) -> int:
22
+ with zstd.open(sinks.raw_out_path, mode="wb") as comp_out:
23
+ try:
24
+ main_proc = subprocess.Popen(
25
+ cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=env, cwd=cwd
26
+ )
27
+ except FileNotFoundError as e:
28
+ raise CommandNotFoundError(cmd[0]) from e
29
+
30
+ return _pump(main_proc, comp_out, mirror_to_stdout=sinks.mirror_to_stdout)
31
+
32
+
33
+ def _pump(
34
+ main_proc: subprocess.Popen[bytes],
35
+ comp_out: zstd.ZstdFile,
36
+ *,
37
+ mirror_to_stdout: bool,
38
+ ) -> int:
39
+ assert main_proc.stdout is not None
40
+ main_stdout = main_proc.stdout
41
+
42
+ try:
43
+ for chunk in iter(lambda: main_stdout.read(io.DEFAULT_BUFFER_SIZE), b""):
44
+ comp_out.write(chunk)
45
+ if mirror_to_stdout:
46
+ sys.stdout.buffer.write(chunk)
47
+ sys.stdout.buffer.flush()
48
+ except BrokenPipeError, KeyboardInterrupt:
49
+ pass
50
+ finally:
51
+ main_stdout.close()
52
+
53
+ return main_proc.wait()
File without changes
@@ -0,0 +1,167 @@
1
+ # SPDX-FileCopyrightText: AISEC Pentesting Team
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import os
6
+ import shlex
7
+ from dataclasses import dataclass
8
+ from pathlib import Path
9
+
10
+ from penrun.artifacts import (
11
+ ArtifactsDir,
12
+ ArtifactsPlaceholders,
13
+ now_iso,
14
+ render_artifactsdir,
15
+ timestamp_now,
16
+ update_latest_symlink,
17
+ )
18
+ from penrun.config import Config
19
+ from penrun.process import PipelineSinks, run_pipeline
20
+
21
+ DEFAULT_TEMPLATE = "{base}/{scriptname}{tagsuffix}/run-{timestamp}"
22
+ DEFAULT_BATCH_TEMPLATE = "{base}/run-{timestamp}/{jobname}{tagsuffix}"
23
+ # Used instead of DEFAULT_BATCH_TEMPLATE when -N/--name-by-arg is given: jobs
24
+ # are named after their input line, so grouping them under the shared
25
+ # run-{timestamp} dir wouldn't tell different commands' batches apart. This
26
+ # nests run-{timestamp} under a directory named after the command instead,
27
+ # mirroring DEFAULT_TEMPLATE's job/datetime layout.
28
+ DEFAULT_BATCH_TEMPLATE_NAMED = "{base}/{scriptname}{tagsuffix}/run-{timestamp}/{jobname}"
29
+
30
+
31
+ def _resolve_executable(cmd: list[str]) -> list[str]:
32
+ exe = cmd[0]
33
+ if os.sep in exe or (os.altsep and os.altsep in exe):
34
+ exe = str(Path(exe).resolve())
35
+ return [exe, *cmd[1:]]
36
+
37
+
38
+ @dataclass(frozen=True, slots=True)
39
+ class RunRequest:
40
+ cmd: list[str]
41
+ tag: str = ""
42
+ artifactsdir_override: Path | None = None
43
+ base_override: Path | None = None
44
+ chdir: bool = False
45
+ # Explicit overrides for the {scriptname}/{jobname} artifacts_template
46
+ # placeholders. None means: derive from cmd[0] (see _program_name).
47
+ scriptname: str | None = None
48
+ jobname: str | None = None
49
+ mirror_to_stdout: bool = True
50
+ # Set (along with `timestamp`) by BatchRunner so every job of one
51
+ # batch shares a single timestamp and picks DEFAULT_BATCH_TEMPLATE.
52
+ batch: bool = False
53
+ # Batch mode picked DEFAULT_BATCH_TEMPLATE_NAMED (jobs nested under the
54
+ # shared scriptname dir) rather than DEFAULT_BATCH_TEMPLATE (jobs named
55
+ # after the plain scriptname) - see -N/--name-by-arg.
56
+ named_jobs: bool = False
57
+ timestamp: str | None = None
58
+ include_env: bool = False
59
+
60
+
61
+ class Runner:
62
+ def __init__(self, config: Config) -> None:
63
+ self._config = config
64
+
65
+ def run(self, request: RunRequest) -> int:
66
+ cmd = request.cmd
67
+
68
+ artifactsdir, latest_dir = self._resolve_paths(request)
69
+ artifacts = ArtifactsDir(artifactsdir, latest_dir)
70
+ artifacts.create()
71
+
72
+ if request.chdir:
73
+ # The child's cwd is about to become artifacts.path, so a
74
+ # relative executable path (e.g. "./scan.sh") must be resolved
75
+ # against *this* process's cwd now, or the child would look
76
+ # for it inside artifacts.path instead.
77
+ cmd = _resolve_executable(cmd)
78
+
79
+ env = os.environ.copy()
80
+ env["PENRUN_COMMAND"] = shlex.join(cmd)
81
+ env["PENRUN_ARTIFACTS"] = str(artifacts.path)
82
+
83
+ return self._run_command(artifacts, cmd, env, request)
84
+
85
+ def _resolve_paths(self, request: RunRequest) -> tuple[Path, Path]:
86
+ if request.artifactsdir_override is not None:
87
+ artifactsdir = self._absolute(request.artifactsdir_override)
88
+ return artifactsdir, artifactsdir.parent
89
+
90
+ jobname = request.jobname or self._program_name(request)
91
+ placeholders = ArtifactsPlaceholders(
92
+ base=str(self._base_dir(request)),
93
+ scriptname=self._program_name(request),
94
+ tag=request.tag,
95
+ fragment=jobname,
96
+ jobname=jobname,
97
+ timestamp=request.timestamp or timestamp_now(),
98
+ )
99
+ template = self._config.artifacts_template or self._default_template(request)
100
+
101
+ artifactsdir = self._absolute(render_artifactsdir(template, placeholders))
102
+ return artifactsdir, artifactsdir.parent
103
+
104
+ def _default_template(self, request: RunRequest) -> str:
105
+ if not request.batch:
106
+ return DEFAULT_TEMPLATE
107
+ if request.named_jobs:
108
+ return DEFAULT_BATCH_TEMPLATE_NAMED
109
+ return DEFAULT_BATCH_TEMPLATE
110
+
111
+ def _program_name(self, request: RunRequest) -> str:
112
+ return request.scriptname or Path(request.cmd[0]).stem
113
+
114
+ def _base_dir(self, request: RunRequest) -> Path:
115
+ base = request.base_override or self._config.artifacts_base or Path.cwd().joinpath("penrun-artifacts")
116
+ return self._absolute(Path(base))
117
+
118
+ def _absolute(self, path: Path) -> Path:
119
+ return path if path.is_absolute() else Path.cwd().joinpath(path)
120
+
121
+ def update_batch_latest(self, request: RunRequest) -> None:
122
+ # Batch jobs share one run-{timestamp} directory (see
123
+ # DEFAULT_BATCH_TEMPLATE / DEFAULT_BATCH_TEMPLATE_NAMED), so a
124
+ # per-job LATEST symlink placed next to it would just have every
125
+ # job in the batch fight over the same link. Instead, batch mode
126
+ # gets a single LATEST pointing at that shared run directory - at
127
+ # the top of the artifacts base normally, or nested under the
128
+ # command's own directory when -N/--name-by-arg groups jobs under
129
+ # DEFAULT_BATCH_TEMPLATE_NAMED instead.
130
+ if request.artifactsdir_override is not None:
131
+ return
132
+ artifactsdir, _ = self._resolve_paths(request)
133
+ run_dir = artifactsdir.parent
134
+ if request.named_jobs and self._config.artifacts_template is None:
135
+ latest_dir = run_dir.parent
136
+ else:
137
+ latest_dir = self._base_dir(request)
138
+ update_latest_symlink(latest_dir.joinpath("LATEST"), run_dir)
139
+
140
+ def _run_command(
141
+ self,
142
+ artifacts: ArtifactsDir,
143
+ cmd: list[str],
144
+ env: dict[str, str],
145
+ request: RunRequest,
146
+ ) -> int:
147
+ if not request.batch:
148
+ artifacts.update_latest_symlink()
149
+
150
+ start = now_iso()
151
+ sinks = PipelineSinks(
152
+ raw_out_path=artifacts.output_path(),
153
+ mirror_to_stdout=request.mirror_to_stdout,
154
+ )
155
+ cwd = artifacts.path if request.chdir else None
156
+ exit_code = run_pipeline(cmd, env, sinks, cwd=cwd)
157
+ end = now_iso()
158
+
159
+ artifacts_env = env if request.include_env else None
160
+ artifacts.write_meta(
161
+ command=cmd,
162
+ start_time=start,
163
+ end_time=end,
164
+ exit_code=exit_code,
165
+ env=artifacts_env,
166
+ )
167
+ return exit_code