diff --git a/SKILLS.md b/SKILLS.md index 917316c680a..8506ff43e2c 100644 --- a/SKILLS.md +++ b/SKILLS.md @@ -72,6 +72,105 @@ CAPE (Config And Payload Extraction) is a malware analysis sandbox derived from * **Logging:** Use `import logging; log = logging.getLogger(__name__)`. Do not use `print()`. * **Exceptions:** Use custom exceptions from `lib/cuckoo/common/exceptions.py` (e.g., `CuckooOperationalError`). +### Local Development Environment + +Dependencies are managed with **Poetry**; the test suite imports `sqlalchemy`, +`django` and friends, so a system Python without the project environment will +fail at collection time in `tests/conftest.py`. + +```bash +poetry install +poetry run pytest tests/ -q +``` + +Use `poetry run ` (or activate the venv) for *every* Python invocation - +`pytest`, `ruff`, `black`, `alembic`, `python utils/...`. + +**Linting.** `ruff` is the fast gate and is expected to be clean: + +```bash +poetry run ruff check +``` + +`black` and `ruff format` are configured with `line-length = 132` in +`pyproject.toml`. Some long-lived modules predate the current settings, so +`black --check` reports findings unrelated to your change. Format only the +lines you touched; reformatting a whole module turns a small fix into an +unreviewable diff. + +**Tests.** Third-party dependencies emit several hundred deprecation warnings +that bury the result; `-p no:warnings` keeps the output readable: + +```bash +poetry run pytest tests/test_something.py -p no:warnings -q +``` + +### Isolated Checkouts for Review and Testing + +Reviewing a PR or reproducing a bug against another branch should never +disturb your working clone - most clones carry uncommitted work, and +`git checkout` / `git stash` in the wrong directory loses it. Use +`utils/agent_worktree.py`, a thin wrapper over `git worktree` that handles the +PR plumbing. It is stdlib-only and repository-agnostic. + +```bash +python utils/agent_worktree.py new --pr 3219 # resolve PR head via gh, fetch, branch, check out +python utils/agent_worktree.py new --branch topic # existing remote branch +python utils/agent_worktree.py new --from origin/master --name scratch + +python utils/agent_worktree.py list # flags: main / managed / dirty / pr#N +python utils/agent_worktree.py path pr3219 # for use in $(...) +python utils/agent_worktree.py update pr3219 # re-fetch after the author pushes +python utils/agent_worktree.py remove pr3219 # also drops the branch it created +python utils/agent_worktree.py cleanup # remove every worktree it created +python utils/agent_worktree.py info # repo, remotes, virtualenv, dirty state +``` + +`new --pr N` asks `gh` which fork the head branch lives in, fetches from the +matching remote (or straight from the fork URL if no remote is configured), +creates a local branch tracking it, and prints the path plus the project +virtualenv. Add `--json` to any command for scripted use. + +Notes: + +* Worktrees default to `~/.cache/agent-worktrees//`; override with + `--base-dir`, `--path`, or `$AGENT_WORKTREE_DIR`. +* Only worktrees created by the tool are ever removed. They are tagged with an + `agent-meta.json` inside the repository's git admin directory, so nothing + extra shows up in `git status`. +* `remove` and `cleanup` refuse to discard uncommitted changes or unpushed + commits; `cleanup` reports what it skipped and why. `--force` overrides. +* Poetry keys virtualenvs by project path, so a fresh worktree has no + environment of its own. Run `poetry run` from your main clone, or use the + interpreter reported by `info`. + +Tests live in `tests/test_agent_worktree.py` and run against throwaway local +repositories - no network, credentials or CAPE configuration required. + +### Codebase Gotchas + +Behaviours that regularly cause silent, hard-to-debug failures: + +* **`lib/cuckoo/common/dictionary.py`** - `Dictionary.__getattr__` returns + `None` for missing keys instead of raising. Consequently + `getattr(section, "key", default)` **never applies its default**; use + `section.get("key", default)`. +* **`lib/cuckoo/common/config.py`** - `_BaseConfig.get(section)` takes exactly + one argument and raises `CuckooOperationalError` for an unknown section, so + `conf.get(name, default)` is a `TypeError`, not a fallback. `Config` is + cached per configuration file by its metaclass, so repeated + `Config("processing")` calls are cheap. +* **`lib/cuckoo/common/integrations/utils.py`** - `run_tool()` returns + **stdout only**. Callers that need stderr must pass + `stderr=subprocess.PIPE` themselves. +* **`lib/cuckoo/common/integrations/file_extra_info_modules/__init__.py`** - + `extractor_ctx()` wraps extractors in `except Exception: log.exception(...)`. + A `TypeError` or `NameError` therefore surfaces as an empty result rather + than a crash; check the log before assuming the logic is wrong. +* Broad `except Exception` around encode/decode helpers hides missing + module-level imports, which then look like a working fallback path. Verify + the import exists. + ### How to Add a Detection Signature Signatures live in `modules/signatures/`. * **Ref:** `docs/book/src/customization/signatures.rst` @@ -204,6 +303,7 @@ If the Python controller is unresponsive, use `py-spy` to inspect the stack trac * **Clean All:** `sudo -u cape poetry run python utils/cleaners.py --clean` (Destructive!) * **Download Signatures:** `sudo -u cape poetry run python utils/community.py -waf` * **Test Rooter:** `sudo python3 utils/rooter.py -g cape -v` +* **Check Out a PR Safely:** `python utils/agent_worktree.py new --pr ` (isolated worktree; your clone is untouched) ### Database Querying (MongoDB) CAPE stores unstructured analysis results in the `analysis` collection. diff --git a/analyzer/windows/modules/auxiliary/watchdownloads.py b/analyzer/windows/modules/auxiliary/watchdownloads.py index c2aefa5ffda..625cd5b0a9e 100644 --- a/analyzer/windows/modules/auxiliary/watchdownloads.py +++ b/analyzer/windows/modules/auxiliary/watchdownloads.py @@ -23,16 +23,33 @@ from watchdog.observers import Observer class MyEventHandler(FileSystemEventHandler): + def __init__(self): + super().__init__() + self._uploaded = {} + def on_any_event(self, event: FileSystemEvent) -> None: - if event.event_type == EVENT_TYPE_DELETED: + if event.is_directory or event.event_type in (EVENT_TYPE_DELETED, "opened"): + return + target_path = getattr(event, "dest_path", None) or event.src_path + if not target_path: return try: - filename = os.path.basename(event.src_path) - if not filename.endswith((".part", "desktop.ini")): - log.info("Monitor uploading %s", filename) - upload_to_host(event.src_path, f"files/{filename}") + filename = os.path.basename(target_path) + if filename.endswith((".part", ".tmp", ".crdownload", "desktop.ini")): + return + if not os.path.exists(target_path): + return + st = os.stat(target_path) + if st.st_size == 0: + return + sig = (st.st_size, st.st_mtime) + if self._uploaded.get(target_path) == sig: + return + log.info("Monitor uploading %s", filename) + upload_to_host(target_path, f"files/{filename}") + self._uploaded[target_path] = sig except Exception as e: - log.exception("Can't upload new file %s to host. %s", event.src_path, str(e)) + log.exception("Can't upload new file %s to host. %s", target_path, str(e)) HAVE_WATCHDOG = True except ImportError as e: diff --git a/changelog.md b/changelog.md index f286fba9744..f843e03bcef 100644 --- a/changelog.md +++ b/changelog.md @@ -2,6 +2,7 @@ * Monitor updates: * Misc enhancements & fixes (see capemon repo for details) + ### [22.08.2026] * Performance & Database Infrastructure: * **psycopg3 Support**: Upgraded the PostgreSQL database connection driver to `psycopg` (v3) for modern async capability and massive performance gains. diff --git a/lib/cuckoo/common/integrations/file_extra_info.py b/lib/cuckoo/common/integrations/file_extra_info.py index b324623b204..4ced2865224 100644 --- a/lib/cuckoo/common/integrations/file_extra_info.py +++ b/lib/cuckoo/common/integrations/file_extra_info.py @@ -253,7 +253,8 @@ def static_file_info( # Below the libmagic "type" already present in data_dictionary. Cached # in the magika integration, so this is a no-op lookup for anything # that already went through File.get_all(). - if processing_conf.magika.enabled and "magika" not in data_dictionary: + magika_cfg = getattr(processing_conf, "magika", None) + if magika_cfg and getattr(magika_cfg, "enabled", False) and "magika" not in data_dictionary: magika_result = magika_info(file_path) if magika_result: data_dictionary["magika"] = magika_result diff --git a/lib/cuckoo/core/resultserver.py b/lib/cuckoo/core/resultserver.py index eced9b407b6..8d4d97b5a7e 100644 --- a/lib/cuckoo/core/resultserver.py +++ b/lib/cuckoo/core/resultserver.py @@ -3,6 +3,7 @@ # See the file 'docs/LICENSE' for copying permission. import errno +import hashlib import json import logging import multiprocessing @@ -10,6 +11,7 @@ import signal import socket import struct +import tempfile from contextlib import suppress from threading import Lock, Thread @@ -179,15 +181,22 @@ def read_newline(self): line, self.buf = self.buf[:pos], self.buf[pos + 1 :] return line - def copy_to_fd(self, fd, max_size=None): + def copy_to_fd(self, fd, max_size=None, hasher=None): if max_size: - fd = WriteLimiter(fd, max_size) - fd.write(self.drain_buffer()) + fd = WriteLimiter(fd, max_size, hasher=hasher) + hasher = None + buf = self.drain_buffer() + if buf: + fd.write(buf) + if hasher is not None: + hasher.update(buf) while True: buf = self.read() if buf == b"": break fd.write(buf) + if hasher is not None: + hasher.update(buf) fd.flush() def discard(self): @@ -201,9 +210,10 @@ def __del__(self): class WriteLimiter: - def __init__(self, fd, remain): + def __init__(self, fd, remain, hasher=None): self.fd = fd self.remain = remain + self.hasher = hasher self.warned = False def write(self, buf): @@ -211,12 +221,18 @@ def write(self, buf): write = min(size, self.remain) try: if write: - self.fd.write(buf[:write]) + chunk = buf[:write] + self.fd.write(chunk) + if self.hasher is not None: + self.hasher.update(chunk) self.remain -= write if size and size != write: if not self.warned: log.warning("Uploaded file length larger than upload_max_size, stopping upload") - self.fd.write(b"... (truncated)") + trunc = b"... (truncated)" + self.fd.write(trunc) + if self.hasher is not None: + self.hasher.update(trunc) self.warned = True except Exception as e: log.debug("Failed to upload file due to '%s'", e) @@ -224,11 +240,56 @@ def write(self, buf): def flush(self): self.fd.flush() + def tell(self): + return self.fd.tell() + def __del__(self): if self.fd: self.fd.close() +# Per-task locks and state for ResultServer file deduplication and versioning +_task_file_state_lock = Lock() +_task_file_locks = {} +_task_sha256_to_path = {} +_task_path_to_sha256 = {} +_task_filelog_seen = {} + + +def _get_task_file_state(task_id): + with _task_file_state_lock: + if task_id not in _task_file_locks: + _task_file_locks[task_id] = Lock() + _task_sha256_to_path[task_id] = {} + _task_path_to_sha256[task_id] = {} + _task_filelog_seen[task_id] = set() + return ( + _task_file_locks[task_id], + _task_sha256_to_path[task_id], + _task_path_to_sha256[task_id], + _task_filelog_seen[task_id], + ) + + +def _cleanup_task_file_state(task_id): + with _task_file_state_lock: + _task_file_locks.pop(task_id, None) + _task_sha256_to_path.pop(task_id, None) + _task_path_to_sha256.pop(task_id, None) + _task_filelog_seen.pop(task_id, None) + + +def _compute_file_sha256(filepath): + h = hashlib.sha256() + try: + with open(filepath, "rb") as f: + while chunk := f.read(65536): + h.update(chunk) + return h.hexdigest() + except OSError: + return None + + class FileUpload(ProtocolHandler): def init(self): self.upload_max_size = cfg.resultserver.upload_max_size @@ -240,6 +301,38 @@ def __del__(self): if self.fd: self.fd.close() + def _write_filelog(self, rel_path, filepath, pids, ppids, metadata, category, seen_set=None): + if rel_path.startswith( + ("shots/", "curtain/", "aux/", "sysmon/", "debugger/", "tlsdump/", "evtx", "htmldump/") + ): + return + + filepath_str = filepath.decode("utf-8", "replace") if filepath else "" + metadata_str = metadata.decode("utf-8", "replace") if metadata else "" + cat_str = category.decode() if category in (b"CAPE", b"files", b"memory", b"procdump") else "" + + entry_key = (rel_path, filepath_str, tuple(pids), tuple(ppids), metadata_str, cat_str) + if seen_set is not None: + if entry_key in seen_set: + return + seen_set.add(entry_key) + + with open(self.filelog, "a") as f: + print( + json.dumps( + { + "path": rel_path, + "filepath": filepath_str, + "pids": pids, + "ppids": ppids, + "metadata": metadata_str, + "category": cat_str, + }, + ensure_ascii=False, + ), + file=f, + ) + def handle(self): # Read until newline for file path, e.g., # shots/0001.jpg or files/9498687557/libcurl-4.dll.bin @@ -261,68 +354,159 @@ def handle(self): else: filepath, pids, ppids, metadata, category, duplicated = None, [], [], b"", b"", False - log.debug("Task #%s: Uploading file %s", self.task_id, dump_path.decode()) - if not duplicated: - file_path = os.path.join(self.storagepath, dump_path.decode()) + rel_dump_path = dump_path.decode("utf-8", "replace") + log.debug("Task #%s: Uploading file %s", self.task_id, rel_dump_path) - try: - if file_path.endswith("_script.log"): - self.fd = open_inclusive(file_path) - elif is_replaceable_result_upload(dump_path) and path_exists(file_path): - # Auxiliary modules (tlsdump, network_etw, sslkeylogfile…) - # upload the SAME dump_path periodically so accumulated - # key / connection data survives an unexpected analysis - # termination. Each upload is a full replacement of the - # prior content — truncate and rewrite rather than failing - # silently with EEXIST. - self.fd = open(file_path, "wb") - else: - # open_exclusive will fail if file_path already exists - self.fd = open_exclusive(file_path) - except OSError as e: - log.debug("File upload error for %s (task #%s)", dump_path, self.task_id) - if e.errno == errno.EEXIST: - raise CuckooOperationalError( - "Task #%s: Analyzer tried to overwrite an existing file: %s" % (self.task_id, file_path) - ) - raise - # ToDo we need Windows path - # filter screens/curtain/sysmon - if not dump_path.startswith( - (b"shots/", b"curtain/", b"aux/", b"sysmon/", b"debugger/", b"tlsdump/", b"evtx", b"htmldump/") - ): - # Append-writes are atomic - with open(self.filelog, "a") as f: - print( - json.dumps( - { - "path": dump_path.decode("utf-8", "replace"), - "filepath": filepath.decode("utf-8", "replace") if filepath else "", - "pids": pids, - "ppids": ppids, - "metadata": metadata.decode("utf-8", "replace"), - "category": category.decode() if category in (b"CAPE", b"files", b"memory", b"procdump") else "", - }, - ensure_ascii=False, - ), - file=f, - ) + task_lock, sha256_map, path_map, seen_set = _get_task_file_state(self.task_id) + + if duplicated: + with task_lock: + self._write_filelog(rel_dump_path, filepath, pids, ppids, metadata, category, seen_set=seen_set) + return - if not duplicated: + file_path = os.path.join(self.storagepath, rel_dump_path) + os.makedirs(os.path.dirname(file_path), exist_ok=True) + + if file_path.endswith("_script.log") or is_replaceable_result_upload(dump_path): + self.fd = open_inclusive(file_path) if file_path.endswith("_script.log") else open(file_path, "wb") + with task_lock: + self._write_filelog(rel_dump_path, filepath, pids, ppids, metadata, category, seen_set=seen_set) self.handler.sock.settimeout(None) try: return self.handler.copy_to_fd(self.fd, self.upload_max_size) except Exception as e: - if self.fd: + log.debug("Task #%s: Failed to upload replaceable log %s due to '%s'", self.task_id, rel_dump_path, e) + return + + # Stream to a temporary file in the target folder while computing SHA256 + hasher = hashlib.sha256() + tmp_fd = tempfile.NamedTemporaryFile( + dir=os.path.dirname(file_path), + prefix=".tmp_upload_", + delete=False, + ) + tmp_path = tmp_fd.name + self.fd = tmp_fd + self.handler.sock.settimeout(None) + try: + self.handler.copy_to_fd(self.fd, self.upload_max_size, hasher=hasher) + except Exception as e: + log.debug( + "Task #%s: Failed to upload file %s due to '%s'", + self.task_id, + rel_dump_path, + e, + ) + self.fd.close() + self.fd = None + with suppress(OSError): + os.unlink(tmp_path) + return + finally: + if self.fd: + self.fd.close() + self.fd = None + + new_sha256 = hasher.hexdigest() + try: + new_size = os.path.getsize(tmp_path) + except OSError: + new_size = 0 + + with task_lock: + # Ensure existing target file on disk is indexed in path_map / sha256_map + if path_exists(file_path) and rel_dump_path not in path_map: + existing_sha = _compute_file_sha256(file_path) + if existing_sha: + path_map[rel_dump_path] = existing_sha + try: + ex_size = os.path.getsize(file_path) + except OSError: + ex_size = 0 + if ex_size > 0 or existing_sha not in sha256_map: + sha256_map.setdefault(existing_sha, rel_dump_path) + + # Check if we already stored a file with this exact SHA256 + existing_same_hash_rel = sha256_map.get(new_sha256) + if existing_same_hash_rel: + existing_same_hash_full = os.path.join(self.storagepath, existing_same_hash_rel) + try: + ex_same_size = os.path.getsize(existing_same_hash_full) if path_exists(existing_same_hash_full) else -1 + except OSError: + ex_same_size = -1 + if ex_same_size < 0 or (ex_same_size == 0 and new_size > 0): + existing_same_hash_rel = None + + if existing_same_hash_rel: + # Duplicate SHA256: discard temp file and map metadata to existing file + with suppress(OSError): + os.unlink(tmp_path) + final_rel_path = existing_same_hash_rel + log.debug( + "Task #%s: Deduplicated uploaded file %s (SHA256 %s matches %s)", + self.task_id, + rel_dump_path, + new_sha256, + final_rel_path, + ) + elif not path_exists(file_path): + # Destination is free and content is unique + os.replace(tmp_path, file_path) + path_map[rel_dump_path] = new_sha256 + sha256_map[new_sha256] = rel_dump_path + final_rel_path = rel_dump_path + else: + # Destination path already exists on disk + try: + existing_size = os.path.getsize(file_path) + except OSError: + existing_size = 0 + + if existing_size == 0 and new_size > 0: + # Replace 0-byte placeholder file in-place + old_sha = path_map.get(rel_dump_path) + if old_sha and sha256_map.get(old_sha) == rel_dump_path: + sha256_map.pop(old_sha, None) + os.replace(tmp_path, file_path) + path_map[rel_dump_path] = new_sha256 + sha256_map[new_sha256] = rel_dump_path + final_rel_path = rel_dump_path log.debug( - "Task #%s: Failed to uploaded file %s of length %s due to '%s'", + "Task #%s: Replaced 0-byte file %s with %d bytes (SHA256 %s)", self.task_id, - dump_path.decode(), - self.fd.tell(), - e, + rel_dump_path, + new_size, + new_sha256, ) + elif new_size == 0 and existing_size > 0: + # Ignore 0-byte upload when a non-empty file already exists + with suppress(OSError): + os.unlink(tmp_path) + final_rel_path = rel_dump_path else: - log.debug("Task #%s: Failed to uploaded file %s due to '%s'", self.task_id, dump_path.decode(), e) + # Different content for same filename -> create versioned file + dirname, basename = os.path.split(rel_dump_path) + stem, ext = os.path.splitext(basename) + version = 1 + while True: + candidate_rel = os.path.join(dirname, f"{stem}_{version}{ext}") if dirname else f"{stem}_{version}{ext}" + candidate_full = os.path.join(self.storagepath, candidate_rel) + if not path_exists(candidate_full): + break + version += 1 + os.replace(tmp_path, candidate_full) + path_map[candidate_rel] = new_sha256 + sha256_map[new_sha256] = candidate_rel + final_rel_path = candidate_rel + log.info( + "Task #%s: Versioned colliding file %s -> %s (SHA256 %s)", + self.task_id, + rel_dump_path, + final_rel_path, + new_sha256, + ) + + self._write_filelog(final_rel_path, filepath, pids, ppids, metadata, category, seen_set=seen_set) class LogHandler(ProtocolHandler): @@ -533,6 +717,7 @@ def del_task(self, task_id, ipaddr): log.debug("Task #%s: Cancel %s", task_id, ctx) ctx.cancel() task_log_stop_force(task_id) + _cleanup_task_file_state(task_id) def create_folders(self): for folder in list(RESULT_UPLOADABLE) + [b"logs"]: @@ -693,6 +878,7 @@ def cancel_task(self, task_id): for ctx in ctxs: ctx.cancel() task_log_stop_force(task_id) + _cleanup_task_file_state(task_id) # Use a SPAWN context (not the default fork) for the per-VM ResultServer worker diff --git a/modules/processing/CAPE.py b/modules/processing/CAPE.py index f89b98b867c..2ec495d5181 100644 --- a/modules/processing/CAPE.py +++ b/modules/processing/CAPE.py @@ -138,7 +138,7 @@ def _metadata_processing(self, metadata, file_info, append_file): if len(metastrings) > 3: file_info["module_path"] = _clean_path(metastrings[2], self.options.replace_patterns) - if "pids" in metadata: + if metadata.get("pids"): file_info["pid"] = metadata["pids"][0] if len(metadata["pids"]) == 1 else ",".join(str(p) for p in metadata["pids"]) if metastrings and metastrings[0] and metastrings[0].isdigit(): @@ -251,7 +251,8 @@ def process_file(self, file_path, append_file, metadata: dict, *, category: str, file_info["type"] = f.get_type() # `file_info` can come straight from the mongo file cache, which may # predate magika being enabled (or a model change). Backfill it. - if processing_conf.magika.enabled and "magika" not in file_info: + magika_cfg = getattr(processing_conf, "magika", None) + if magika_cfg and getattr(magika_cfg, "enabled", False) and "magika" not in file_info and hasattr(f, "get_magika"): magika_result = f.get_magika() if magika_result: file_info["magika"] = magika_result @@ -301,12 +302,20 @@ def process_file(self, file_path, append_file, metadata: dict, *, category: str, if category == "dropped": file_info.update(metadata.get(file_info["path"][0], {})) file_info["guest_paths"] = list( - {_clean_path(path.get("filepath", ""), self.options.replace_patterns) for path in metadata.get(file_path, [])} + dict.fromkeys( + _clean_path(path.get("filepath", ""), self.options.replace_patterns) + for path in metadata.get(file_path, []) + if path.get("filepath") + ) ) if not file_info["guest_paths"] and category == "dropped" and "CAPE" not in metadata.get("filepath", ""): file_info["guest_paths"] = [_clean_path(metadata.get("filepath", ""), self.options.replace_patterns)] file_info["name"] = list( - {path.get("filepath", "").rsplit("\\", 1)[-1] for path in metadata.get(file_path, [])} + dict.fromkeys( + path.get("filepath", "").rsplit("\\", 1)[-1] + for path in metadata.get(file_path, []) + if path.get("filepath") + ) ) or [metadata.get("filepath", "").rsplit("\\", 1)[-1]] if category == "dropped": with suppress(UnicodeDecodeError): @@ -434,12 +443,22 @@ def run(self): continue filepath = os.path.join(self.analysis_path, entry["path"]) - meta[filepath] = { - "pids": entry.get("pids"), - "ppids": entry.get("ppids"), - "filepath": entry.get("filepath", ""), - "metadata": entry.get("metadata", {}), - } + if filepath in meta: + for p in entry.get("pids") or []: + if p not in meta[filepath]["pids"]: + meta[filepath]["pids"].append(p) + for p in entry.get("ppids") or []: + if p not in meta[filepath]["ppids"]: + meta[filepath]["ppids"].append(p) + meta[filepath][filepath].append(entry) + else: + meta[filepath] = { + "pids": list(entry.get("pids") or []), + "ppids": list(entry.get("ppids") or []), + "filepath": entry.get("filepath", ""), + "metadata": entry.get("metadata", {}), + filepath: [entry], + } # Pre-scan ClamAV in parallel for every file we're about to process. # The sequential single-thread `allmatchscan` over 10-20 dropped / @@ -470,7 +489,8 @@ def run(self): # Same lifecycle contract as the clamav cache: drop per-path magika # results at the task boundary so a long-lived worker can't serve a # stale prediction for a path that has been reused by another task. - if processing_conf.magika.enabled: + magika_cfg = getattr(processing_conf, "magika", None) + if magika_cfg and getattr(magika_cfg, "enabled", False): try: from lib.cuckoo.common.integrations.magika import clear_magika_cache diff --git a/tests/test_agent_worktree.py b/tests/test_agent_worktree.py new file mode 100644 index 00000000000..6f7d5ddbad4 --- /dev/null +++ b/tests/test_agent_worktree.py @@ -0,0 +1,317 @@ +# This file is part of CAPE Sandbox - https://github.com/kevoreilly/CAPEv2 +# See the file 'docs/LICENSE' for copying permission. + +"""Tests for utils/agent_worktree.py. + +Everything runs against throwaway git repositories created in a tmp_path, so +the suite needs no network access, no GitHub credentials and no CAPE +configuration. The ``--pr`` code path is the only part not covered here +because it requires the ``gh`` CLI and a live API. +""" + +import json +import os +import subprocess +import sys + +import pytest + +sys.path.append(os.path.join(os.path.dirname(__file__), "..")) + +from utils import agent_worktree # noqa: E402 + + +def _git(repo, *args): + subprocess.run( + ["git", "-C", str(repo), *args], + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + universal_newlines=True, + ) + + +def _commit(repo, name, content="x"): + (repo / name).write_text(content) + _git(repo, "add", name) + _git(repo, "commit", "-m", f"add {name}") + + +@pytest.fixture +def origin(tmp_path): + """A bare repository standing in for the remote.""" + path = tmp_path / "origin.git" + path.mkdir() + _git(path, "init", "--bare", "--initial-branch=master", ".") + return path + + +@pytest.fixture +def repo(tmp_path, origin): + """A clone with one commit on master and a pushed feature branch.""" + path = tmp_path / "clone" + path.mkdir() + _git(path, "init", "--initial-branch=master", ".") + _git(path, "config", "user.email", "test@example.com") + _git(path, "config", "user.name", "test") + _git(path, "remote", "add", "origin", str(origin)) + _commit(path, "README") + _git(path, "push", "-u", "origin", "master") + + _git(path, "checkout", "-b", "feature") + _commit(path, "feature.txt") + _git(path, "push", "-u", "origin", "feature") + _git(path, "checkout", "master") + # Delete the local copy so --branch has to fetch it back. + _git(path, "branch", "-D", "feature") + return path + + +def run(repo, base_dir, *args): + """Invoke the CLI in-process and return parsed JSON output.""" + argv = ["--repo", str(repo), "--base-dir", str(base_dir), "--json", *args] + rc = agent_worktree.main(argv) + assert rc == 0 + return argv + + +def call(capsys, repo, base_dir, *args): + run(repo, base_dir, *args) + return json.loads(capsys.readouterr().out) + + +def test_info_reports_repo_and_remotes(capsys, repo, tmp_path): + data = call(capsys, repo, tmp_path / "wt", "info") + assert data["repo"] == str(repo) + assert "origin" in data["remotes"] + assert data["dirty"] is False + + +def test_new_from_ref_creates_worktree(capsys, repo, tmp_path): + base = tmp_path / "wt" + data = call(capsys, repo, base, "new", "--from", "origin/master", "--name", "scratch") + + assert os.path.isdir(data["path"]) + assert data["path"].startswith(str(base)) + assert data["branch"] == "agent/scratch" + # The source clone still sits on master and is untouched. + head = subprocess.run( + ["git", "-C", str(repo), "rev-parse", "--abbrev-ref", "HEAD"], + stdout=subprocess.PIPE, + universal_newlines=True, + check=True, + ).stdout.strip() + assert head == "master" + + +def test_new_branch_fetches_and_tracks_upstream(capsys, repo, tmp_path): + data = call(capsys, repo, tmp_path / "wt", "new", "--branch", "feature") + + assert data["branch"] == "feature" + assert data["upstream"] == "origin/feature" + assert os.path.isfile(os.path.join(data["path"], "feature.txt")) + + +def test_branch_collision_gets_a_suffix(capsys, repo, tmp_path): + base = tmp_path / "wt" + first = call(capsys, repo, base, "new", "--from", "origin/master", "--name", "one", "--branch-name", "dup") + second = call(capsys, repo, base, "new", "--from", "origin/master", "--name", "two", "--branch-name", "dup") + + assert first["branch"] == "dup" + assert second["branch"] == "dup-2" + assert any("was taken" in w for w in second["warnings"]) + + +def test_list_marks_main_and_managed(capsys, repo, tmp_path): + base = tmp_path / "wt" + created = call(capsys, repo, base, "new", "--from", "origin/master", "--name", "scratch") + data = call(capsys, repo, base, "list") + + by_path = {row["path"]: row for row in data["worktrees"]} + assert by_path[str(repo)]["main"] is True + assert by_path[str(repo)]["managed"] is False + assert by_path[created["path"]]["managed"] is True + assert by_path[created["path"]]["name"] == "scratch" + + +def test_path_resolves_name_and_branch(capsys, repo, tmp_path): + base = tmp_path / "wt" + created = call(capsys, repo, base, "new", "--from", "origin/master", "--name", "scratch") + + assert call(capsys, repo, base, "path", "scratch")["path"] == created["path"] + assert call(capsys, repo, base, "path", "agent/scratch")["path"] == created["path"] + + +def test_path_rejects_unknown_target(repo): + # Refusals are reported as a non-zero exit code, not a traceback. + assert agent_worktree.main(["--repo", str(repo), "--json", "path", "nope"]) == 1 + + +def test_remove_deletes_worktree_and_branch(capsys, repo, tmp_path): + base = tmp_path / "wt" + created = call(capsys, repo, base, "new", "--from", "origin/master", "--name", "scratch") + + data = call(capsys, repo, base, "remove", "scratch") + assert data["branch_deleted"] is True + assert not os.path.isdir(created["path"]) + assert not agent_worktree.branch_exists(str(repo), "agent/scratch") + + +def test_remove_refuses_dirty_worktree(capsys, repo, tmp_path): + base = tmp_path / "wt" + created = call(capsys, repo, base, "new", "--from", "origin/master", "--name", "scratch") + open(os.path.join(created["path"], "junk"), "w").close() + + assert agent_worktree.main(["--repo", str(repo), "--base-dir", str(base), "remove", "scratch"]) == 1 + assert os.path.isdir(created["path"]) + + # --force overrides. + assert agent_worktree.main(["--repo", str(repo), "--base-dir", str(base), "remove", "scratch", "--force"]) == 0 + assert not os.path.isdir(created["path"]) + + +def test_remove_refuses_unpushed_commits(capsys, repo, tmp_path): + base = tmp_path / "wt" + created = call(capsys, repo, base, "new", "--branch", "feature") + _git(created["path"], "config", "user.email", "test@example.com") + _git(created["path"], "config", "user.name", "test") + _git(created["path"], "commit", "--allow-empty", "-m", "local only") + + assert agent_worktree.main(["--repo", str(repo), "--base-dir", str(base), "remove", "feature"]) == 1 + assert os.path.isdir(created["path"]) + + +def test_remove_allows_work_pushed_to_another_remote(capsys, repo, tmp_path): + """Tracking the base branch while pushing elsewhere is not 'unpushed'. + + A review branch typically tracks the branch it will merge into while its + commits are pushed to a fork, so the guard has to look at every remote + rather than only the tracked upstream. + """ + fork = tmp_path / "fork.git" + fork.mkdir() + _git(fork, "init", "--bare", "--initial-branch=master", ".") + _git(repo, "remote", "add", "fork", str(fork)) + + base = tmp_path / "wt" + created = call(capsys, repo, base, "new", "--from", "origin/master", "--name", "scratch") + _git(created["path"], "config", "user.email", "test@example.com") + _git(created["path"], "config", "user.name", "test") + _git(created["path"], "commit", "--allow-empty", "-m", "review fix") + # Published to the fork, while the branch still tracks origin/master. + _git(created["path"], "push", "fork", "HEAD:review-branch") + _git(created["path"], "branch", "--set-upstream-to", "origin/master") + + assert agent_worktree.unpushed(created["path"]) is False + assert agent_worktree.main(["--repo", str(repo), "--base-dir", str(base), "remove", "scratch"]) == 0 + assert not os.path.isdir(created["path"]) + + +def test_remove_refuses_main_worktree(repo): + assert agent_worktree.main(["--repo", str(repo), "remove", str(repo)]) == 1 + assert os.path.isdir(repo) + + +def test_update_resets_to_upstream(capsys, repo, tmp_path, origin): + base = tmp_path / "wt" + created = call(capsys, repo, base, "new", "--branch", "feature") + before = created["head"] + + # Advance the remote branch from a second clone. + other = tmp_path / "other" + other.mkdir() + _git(other, "clone", str(origin), ".") + _git(other, "config", "user.email", "test@example.com") + _git(other, "config", "user.name", "test") + _git(other, "checkout", "feature") + _commit(other, "newer.txt") + _git(other, "push", "origin", "feature") + + data = call(capsys, repo, base, "update", "feature") + assert data["after"] != before[:12] + assert os.path.isfile(os.path.join(created["path"], "newer.txt")) + + +def test_cleanup_skips_unsafe_and_removes_the_rest(capsys, repo, tmp_path): + base = tmp_path / "wt" + safe = call(capsys, repo, base, "new", "--from", "origin/master", "--name", "safe") + dirty = call(capsys, repo, base, "new", "--from", "origin/master", "--name", "dirty") + open(os.path.join(dirty["path"], "junk"), "w").close() + + data = call(capsys, repo, base, "cleanup") + assert [item["path"] for item in data["removed"]] == [safe["path"]] + assert [item["path"] for item in data["skipped"]] == [dirty["path"]] + assert data["skipped"][0]["reason"] == "uncommitted changes" + + data = call(capsys, repo, base, "cleanup", "--force") + assert [item["path"] for item in data["removed"]] == [dirty["path"]] + assert not os.path.isdir(dirty["path"]) + + +def test_cleanup_ignores_unmanaged_worktrees(capsys, repo, tmp_path): + base = tmp_path / "wt" + manual = tmp_path / "manual" + _git(repo, "worktree", "add", "--detach", str(manual), "origin/master") + + call(capsys, repo, base, "cleanup", "--force") + assert os.path.isdir(manual) + + +def test_metadata_lives_outside_the_working_tree(capsys, repo, tmp_path): + created = call(capsys, repo, tmp_path / "wt", "new", "--from", "origin/master", "--name", "scratch") + + status = subprocess.run( + ["git", "-C", created["path"], "status", "--porcelain"], + stdout=subprocess.PIPE, + universal_newlines=True, + check=True, + ).stdout + assert status.strip() == "" + assert agent_worktree.read_meta(created["path"])["name"] == "scratch" + + +def test_global_flags_accepted_on_either_side(capsys, repo, tmp_path): + assert agent_worktree.main(["--json", "--repo", str(repo), "list"]) == 0 + prefix = json.loads(capsys.readouterr().out) + assert agent_worktree.main(["list", "--json", "--repo", str(repo)]) == 0 + suffix = json.loads(capsys.readouterr().out) + assert prefix == suffix + + +def test_parse_slug_handles_common_url_forms(): + assert agent_worktree.parse_slug("git@github.com:kevoreilly/CAPEv2.git") == ("kevoreilly", "CAPEv2") + assert agent_worktree.parse_slug("https://github.com/kevoreilly/CAPEv2") == ("kevoreilly", "CAPEv2") + assert agent_worktree.parse_slug("ssh://git@github.com/kevoreilly/CAPEv2.git") == ("kevoreilly", "CAPEv2") + assert agent_worktree.parse_slug("/some/local/path") is None + + +def test_repo_slug_prefers_upstream_over_origin(repo): + # Fork layout: origin is the contributor's fork, upstream is canonical. + _git(repo, "remote", "set-url", "origin", "git@github.com:contributor/project.git") + _git(repo, "remote", "add", "upstream", "git@github.com:canonical/project.git") + assert agent_worktree.repo_slug(str(repo)) == "canonical/project" + + +def test_new_branch_falls_back_to_other_remotes(capsys, repo, tmp_path, origin): + """A branch that exists only on upstream is still found.""" + second = tmp_path / "second.git" + second.mkdir() + _git(second, "init", "--bare", "--initial-branch=master", ".") + _git(repo, "remote", "add", "upstream", str(second)) + + # Publish a branch to upstream only, then drop every local trace of it. + _git(repo, "checkout", "-b", "upstream-only") + _commit(repo, "upstream_only.txt") + _git(repo, "push", "upstream", "upstream-only") + _git(repo, "checkout", "master") + _git(repo, "branch", "-D", "upstream-only") + + data = call(capsys, repo, tmp_path / "wt", "new", "--branch", "upstream-only") + assert data["upstream"] == "upstream/upstream-only" + assert any("not on origin" in w for w in data["warnings"]) + assert os.path.isfile(os.path.join(data["path"], "upstream_only.txt")) + + +def test_new_branch_reports_a_missing_branch(repo, tmp_path): + rc = agent_worktree.main(["--repo", str(repo), "--base-dir", str(tmp_path / "wt"), "new", "--branch", "nope"]) + assert rc == 1 diff --git a/tests/test_resultserver_dedup.py b/tests/test_resultserver_dedup.py new file mode 100644 index 00000000000..eabb249d611 --- /dev/null +++ b/tests/test_resultserver_dedup.py @@ -0,0 +1,369 @@ +import io +import json +import os +import sys +import tempfile +from unittest.mock import MagicMock + +import types + +class _MockProtocolHandler: + def __init__(self, task_id, ctx, version=None): + self.task_id = task_id + self.handler = ctx + self.fd = None + self.version = version + + def __enter__(self): + self.init() + + def __exit__(self, type, value, traceback): + self.close() + + def close(self): + if self.fd: + self.fd.close() + self.fd = None + + def handle(self): + raise NotImplementedError + + +for _mod_name in ( + "gevent", + "gevent.pool", + "gevent.server", + "gevent.socket", + "gevent.thread", + "gevent.monkey", + "gevent.event", + "gevent.lock", + "pytz", + "pebble", + "dns", + "dns.resolver", + "pefile", + "lib.cuckoo.core.database", + "lib.cuckoo.core.data.db_common", + "lib.cuckoo.core.data.task", + "lib.cuckoo.core.data.machines", +): + try: + __import__(_mod_name) + except Exception: + _m = types.ModuleType(_mod_name) + _m.__getattr__ = lambda name: MagicMock() + sys.modules[_mod_name] = _m + +try: + import lib.cuckoo.common.abstracts # noqa: F401 +except Exception: + abstracts_mod = types.ModuleType("lib.cuckoo.common.abstracts") + abstracts_mod.ProtocolHandler = _MockProtocolHandler + abstracts_mod.Processing = MagicMock + sys.modules["lib.cuckoo.common.abstracts"] = abstracts_mod + +from lib.cuckoo.core.resultserver import FileUpload, HandlerContext + + +class DummyHandlerContext(HandlerContext): + def __init__(self, task_id, storagepath, lines, payload): + self.task_id = task_id + self.storagepath = storagepath + self._lines = [line.encode("utf-8") if isinstance(line, str) else line for line in lines] + self._payload = io.BytesIO(payload) + self.buf = b"" + self.sock = MagicMock() + + def read_newline(self): + if not self._lines: + return b"" + return self._lines.pop(0) + + def copy_to_fd(self, fd, max_size=None, hasher=None): + data = self._payload.read() + fd.write(data) + if hasher is not None: + hasher.update(data) + + +def _upload_file( + task_id, + storagepath, + dump_path, + guest_path, + payload, + pids="100", + ppids="50", + metadata="", + category="files", + duplicated="0", +): + lines = [dump_path, guest_path, pids, ppids, metadata, category, duplicated] + ctx = DummyHandlerContext(task_id, storagepath, lines, payload) + handler = FileUpload(task_id, ctx, version=2) + handler.init() + handler.handle() + + +def test_resultserver_dedup_and_versioning(): + task_id = 999991 + with tempfile.TemporaryDirectory() as storagepath: + payload_a = b"hello world payload A" + payload_b = b"different content payload B" + + # 1. First upload of test.bin + _upload_file( + task_id, + storagepath, + "files/test.bin", + r"C:\Users\admin\Downloads\test.bin", + payload_a, + pids="1001", + ) + + file_a_path = os.path.join(storagepath, "files/test.bin") + assert os.path.exists(file_a_path) + with open(file_a_path, "rb") as f: + assert f.read() == payload_a + + # 2. Duplicate upload of same content under colliding filename (e.g. browser retry test(1).bin) + _upload_file( + task_id, + storagepath, + "files/test(1).bin", + r"C:\Users\admin\Downloads\test(1).bin", + payload_a, + pids="1002", + ) + + # test(1).bin should NOT exist on disk; it should be deduplicated to files/test.bin + assert not os.path.exists(os.path.join(storagepath, "files/test(1).bin")) + + # 3. Duplicate upload of same content under exact same filename + _upload_file( + task_id, + storagepath, + "files/test.bin", + r"C:\Users\admin\Downloads\test.bin", + payload_a, + pids="1001", + ) + assert not os.path.exists(os.path.join(storagepath, "files/test_1.bin")) + + # 4. Different content uploaded to same filename files/test.bin -> should version to files/test_1.bin + _upload_file( + task_id, + storagepath, + "files/test.bin", + r"C:\Users\admin\Downloads\test.bin", + payload_b, + pids="1003", + ) + + versioned_path = os.path.join(storagepath, "files/test_1.bin") + assert os.path.exists(versioned_path) + with open(versioned_path, "rb") as f: + assert f.read() == payload_b + + # Verify files.json entries + files_json = os.path.join(storagepath, "files.json") + assert os.path.exists(files_json) + with open(files_json, "rb") as f: + records = [json.loads(line) for line in f] + + # Expected records: + # - test.bin (pid 1001) -> path: files/test.bin + # - test(1).bin (pid 1002) -> path: files/test.bin (deduplicated!) + # Note: identical repeat (test.bin with pid 1001) is deduplicated from files.json + # - test.bin (pid 1003, payload B) -> path: files/test_1.bin + assert len(records) == 3 + assert records[0]["path"] == "files/test.bin" + assert records[0]["filepath"] == r"C:\Users\admin\Downloads\test.bin" + assert records[1]["path"] == "files/test.bin" + assert records[1]["filepath"] == r"C:\Users\admin\Downloads\test(1).bin" + assert records[2]["path"] == "files/test_1.bin" + assert records[2]["filepath"] == r"C:\Users\admin\Downloads\test.bin" + + +def test_resultserver_zero_byte_placeholder_replacement(): + task_id = 999992 + with tempfile.TemporaryDirectory() as storagepath: + payload_real = b"actual APK content bytes" + + # 1. Initial 0-byte upload (e.g. empty file creation event) + _upload_file( + task_id, + storagepath, + "files/app.apk", + r"C:\Users\admin\Downloads\app.apk", + b"", + pids="2001", + ) + target_path = os.path.join(storagepath, "files/app.apk") + assert os.path.exists(target_path) + assert os.path.getsize(target_path) == 0 + + # 2. Full file content uploaded for same path -> replaces 0-byte file in-place (no app_1.apk) + _upload_file( + task_id, + storagepath, + "files/app.apk", + r"C:\Users\admin\Downloads\app.apk", + payload_real, + pids="2001", + ) + assert os.path.exists(target_path) + assert os.path.getsize(target_path) == len(payload_real) + assert not os.path.exists(os.path.join(storagepath, "files/app_1.apk")) + + # 3. Subsequent 0-byte upload -> ignored, does not overwrite real file + _upload_file( + task_id, + storagepath, + "files/app.apk", + r"C:\Users\admin\Downloads\app.apk", + b"", + pids="2002", + ) + assert os.path.getsize(target_path) == len(payload_real) + + +def test_watchdownloads_event_handler_moved_and_dedup(monkeypatch): + import importlib.util + + os.environ.setdefault("HOMEPATH", "/tmp") + + # Mock guest-only modules and watchdog if not installed + mock_abstracts = types.ModuleType("lib.common.abstracts") + mock_abstracts.Auxiliary = type("Auxiliary", (), {}) + mock_results = types.ModuleType("lib.common.results") + uploaded = [] + mock_results.upload_to_host = lambda src, dst: uploaded.append((src, dst)) + + mock_wd_events = types.ModuleType("watchdog.events") + mock_wd_events.EVENT_TYPE_DELETED = "deleted" + mock_wd_events.FileSystemEvent = object + mock_wd_events.FileSystemEventHandler = object + mock_wd_observers = types.ModuleType("watchdog.observers") + mock_wd_observers.Observer = MagicMock() + + monkeypatch.setitem(sys.modules, "lib.common.abstracts", mock_abstracts) + monkeypatch.setitem(sys.modules, "lib.common.results", mock_results) + monkeypatch.setitem(sys.modules, "watchdog.events", mock_wd_events) + monkeypatch.setitem(sys.modules, "watchdog.observers", mock_wd_observers) + + spec = importlib.util.spec_from_file_location( + "watchdownloads", + os.path.join(os.path.dirname(__file__), "../analyzer/windows/modules/auxiliary/watchdownloads.py"), + ) + watchdownloads = importlib.util.module_from_spec(spec) + spec.loader.exec_module(watchdownloads) + + handler = watchdownloads.MyEventHandler() + + with tempfile.TemporaryDirectory() as tmpdir: + part_path = os.path.join(tmpdir, "sample.apk.part") + final_path = os.path.join(tmpdir, "sample.apk") + + # 1. Create empty placeholder -> should be ignored + with open(final_path, "wb") as f: + pass + ev_empty = MagicMock() + ev_empty.is_directory = False + ev_empty.event_type = "created" + ev_empty.src_path = final_path + ev_empty.dest_path = None + handler.on_any_event(ev_empty) + assert len(uploaded) == 0 + + # 2. Write payload and simulate FileMovedEvent (.part -> .apk) + with open(final_path, "wb") as f: + f.write(b"APK content") + ev_moved = MagicMock() + ev_moved.is_directory = False + ev_moved.event_type = "moved" + ev_moved.src_path = part_path + ev_moved.dest_path = final_path + handler.on_any_event(ev_moved) + assert len(uploaded) == 1 + assert uploaded[0] == (final_path, "files/sample.apk") + + # 3. Subsequent modified/closed event with same size & mtime -> should be deduplicated + ev_mod = MagicMock() + ev_mod.is_directory = False + ev_mod.event_type = "modified" + ev_mod.src_path = final_path + ev_mod.dest_path = None + handler.on_any_event(ev_mod) + assert len(uploaded) == 1 + + +def test_cape_processing_without_magika_config(monkeypatch): + """Verify CAPE processing module handles deduplicated dropped files when [magika] config is absent.""" + from lib.cuckoo.common.config import Config + import modules.processing.CAPE as cape_mod + import lib.cuckoo.common.integrations.file_extra_info as fei_mod + + proc_cfg = Config("processing") + if hasattr(proc_cfg, "magika"): + delattr(proc_cfg, "magika") + monkeypatch.setattr(cape_mod, "processing_conf", proc_cfg) + monkeypatch.setattr(fei_mod, "processing_conf", proc_cfg) + monkeypatch.setattr(cape_mod, "static_file_info", lambda *a, **kw: None) + monkeypatch.setattr(cape_mod.File, "get_yara", lambda self, category=None: []) + monkeypatch.setattr("lib.cuckoo.common.objects.get_clamav", lambda path: []) + + with tempfile.TemporaryDirectory() as tmpdir: + files_dir = os.path.join(tmpdir, "files") + os.makedirs(files_dir) + sample_path = os.path.join(files_dir, "Reader_en_install.exe") + with open(sample_path, "wb") as f: + f.write(b"MZ\x90\x00" + b"A" * 1024) + + files_json = os.path.join(tmpdir, "files.json") + with open(files_json, "w") as f: + f.write( + json.dumps( + { + "path": "files/Reader_en_install.exe", + "filepath": "C:\\Users\\Bruno\\Downloads\\Reader_en_install.exe", + "pids": [3540], + "ppids": [1412], + "metadata": "", + "category": "files", + } + ) + + "\n" + ) + f.write( + json.dumps( + { + "path": "files/Reader_en_install.exe", + "filepath": "C:\\Users\\Bruno\\Downloads\\Reader_en_install(1).exe", + "pids": [3540], + "ppids": [1412], + "metadata": "", + "category": "files", + } + ) + + "\n" + ) + + results = {} + cape_proc = cape_mod.CAPE(results) + cape_proc.set_path(tmpdir) + cape_proc.set_task({"id": 1418, "category": "url", "options": ""}) + cape_proc.set_options(proc_cfg.CAPE) + cape_proc.run() + + assert "dropped" in results + assert len(results["dropped"]) == 1 + dropped = results["dropped"][0] + assert dropped["name"] == ["Reader_en_install.exe", "Reader_en_install(1).exe"] + assert dropped["guest_paths"] == [ + "C:\\Users\\Bruno\\Downloads\\Reader_en_install.exe", + "C:\\Users\\Bruno\\Downloads\\Reader_en_install(1).exe", + ] + + diff --git a/utils/agent_worktree.py b/utils/agent_worktree.py new file mode 100755 index 00000000000..1debda788db --- /dev/null +++ b/utils/agent_worktree.py @@ -0,0 +1,789 @@ +#!/usr/bin/env python3 + +# This file is part of CAPE Sandbox - https://github.com/kevoreilly/CAPEv2 +# See the file 'docs/LICENSE' for copying permission. + +"""agent_worktree.py - create and manage disposable git worktrees. + +Why this exists +--------------- +Reviewing a pull request, bisecting a regression or running the test suite +against someone else's branch all need a checkout that is *not* your working +clone. Most contributors' clones carry uncommitted work, and automation +(CI helpers, coding agents, review bots) has no safe way to know that, so a +stray ``git checkout`` or ``git stash`` loses changes. + +``git worktree`` is the right primitive, but the surrounding plumbing - +resolving which fork a PR's head branch lives in, fetching it under a sane +name, setting upstream, and tearing everything down afterwards - is tedious +and easy to get wrong. This wrapper makes the safe path the short path. + +It is repository-agnostic; nothing in it is CAPE-specific. + +Usage +----- + agent_worktree.py new --pr 3219 + agent_worktree.py new --branch some-topic-branch + agent_worktree.py new --name scratch --from origin/master + agent_worktree.py list + agent_worktree.py path pr3219 + agent_worktree.py update pr3219 + agent_worktree.py info + agent_worktree.py remove pr3219 + agent_worktree.py cleanup --all + +Every command accepts ``--repo PATH`` (defaults to the repository containing +the current directory) and ``--json`` for machine-readable output. + +Design guarantees +----------------- +* The main working tree is never checked out, reset, stashed, or cleaned. +* Only worktrees created by this tool are ever removed (they carry a metadata + marker inside the repo's git admin directory). +* Removal refuses to drop uncommitted work unless ``--force`` is given. +""" + +from __future__ import annotations + +import argparse +import datetime as _dt +import glob +import json +import os +import re +import shutil +import subprocess +import sys +from typing import Any, Dict, List, Optional, Tuple + +DEFAULT_BASE_DIR = os.environ.get("AGENT_WORKTREE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "agent-worktrees")) +META_NAME = "agent-meta.json" +FETCH_NS = "refs/agent-worktree" + + +# --------------------------------------------------------------------------- # +# process helpers +# --------------------------------------------------------------------------- # +class CommandError(RuntimeError): + def __init__(self, cmd: List[str], returncode: int, stderr: str): + self.cmd = cmd + self.returncode = returncode + self.stderr = stderr + super().__init__(f"{' '.join(cmd)} exited {returncode}: {stderr.strip()}") + + +def run(cmd: List[str], cwd: Optional[str] = None, check: bool = True) -> str: + proc = subprocess.run( + cmd, + cwd=cwd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + universal_newlines=True, + ) + if check and proc.returncode != 0: + raise CommandError(cmd, proc.returncode, proc.stderr) + return proc.stdout + + +def git(repo: Optional[str], *args: str, check: bool = True) -> str: + cmd = ["git"] + if repo: + cmd += ["-C", repo] + cmd += list(args) + return run(cmd, check=check) + + +# --------------------------------------------------------------------------- # +# repository discovery +# --------------------------------------------------------------------------- # +def main_worktree(start: Optional[str] = None) -> str: + """Return the main working tree, even when called from inside a worktree.""" + start = start or os.getcwd() + try: + out = git(start, "worktree", "list", "--porcelain") + except CommandError as exc: + raise SystemExit(f"not inside a git repository ({start}): {exc.stderr.strip()}") + for line in out.splitlines(): + if line.startswith("worktree "): + return line.split(" ", 1)[1].strip() + raise SystemExit(f"could not determine main worktree for {start}") + + +def git_common_dir(repo: str) -> str: + out = git(repo, "rev-parse", "--git-common-dir").strip() + return out if os.path.isabs(out) else os.path.abspath(os.path.join(repo, out)) + + +def remotes(repo: str) -> Dict[str, str]: + result: Dict[str, str] = {} + for line in git(repo, "remote", "-v").splitlines(): + parts = line.split() + if len(parts) >= 2: + result.setdefault(parts[0], parts[1]) + return result + + +_URL_RE = re.compile( + r"(?:git@|ssh://git@|https://|git://)" r"(?P[^/:]+)[/:]" r"(?P[^/]+)/" r"(?P[^/]+?)(?:\.git)?/?$" +) + + +def parse_slug(url: str) -> Optional[Tuple[str, str]]: + m = _URL_RE.match(url.strip()) + if not m: + return None + return m.group("owner"), m.group("name") + + +def repo_slug(repo: str) -> Optional[str]: + """Best guess at the canonical owner/repo for this checkout. + + `upstream` wins over `origin`: in a fork workflow `origin` is the + contributor's own fork, while pull requests live on the canonical + repository that `upstream` points at. Clones made directly from the + canonical repository have no `upstream`, so `origin` is correct there. + """ + rem = remotes(repo) + for name in ("upstream", "origin"): + if name in rem: + slug = parse_slug(rem[name]) + if slug: + return f"{slug[0]}/{slug[1]}" + for url in rem.values(): + slug = parse_slug(url) + if slug: + return f"{slug[0]}/{slug[1]}" + return None + + +def remote_for(repo: str, owner: str, name: str) -> Optional[str]: + want = (owner.lower(), name.lower()) + for remote, url in remotes(repo).items(): + slug = parse_slug(url) + if slug and (slug[0].lower(), slug[1].lower()) == want: + return remote + return None + + +# --------------------------------------------------------------------------- # +# worktree / branch bookkeeping +# --------------------------------------------------------------------------- # +def worktrees(repo: str) -> List[Dict[str, Any]]: + entries: List[Dict[str, Any]] = [] + current: Dict[str, Any] = {} + for line in git(repo, "worktree", "list", "--porcelain").splitlines(): + if not line.strip(): + if current: + entries.append(current) + current = {} + continue + if " " in line: + key, value = line.split(" ", 1) + else: + key, value = line, True + if key == "worktree": + if current: + entries.append(current) + current = {"path": value} + elif key == "branch": + current["branch"] = value.replace("refs/heads/", "", 1) + elif key == "HEAD": + current["head"] = value + elif key == "detached": + current["detached"] = True + elif key in ("bare", "locked", "prunable"): + current[key] = value + if current: + entries.append(current) + return entries + + +def checked_out_branches(repo: str) -> Dict[str, str]: + return {w["branch"]: w["path"] for w in worktrees(repo) if w.get("branch")} + + +def branch_exists(repo: str, branch: str) -> bool: + return ( + subprocess.run( + ["git", "-C", repo, "show-ref", "--verify", "--quiet", f"refs/heads/{branch}"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ).returncode + == 0 + ) + + +def admin_dir(worktree_path: str) -> Optional[str]: + """Resolve /worktrees/ for a linked worktree.""" + dot_git = os.path.join(worktree_path, ".git") + if not os.path.isfile(dot_git): + return None + with open(dot_git, "r", encoding="utf-8") as fh: + content = fh.read().strip() + if not content.startswith("gitdir:"): + return None + path = content.split(":", 1)[1].strip() + return path if os.path.isdir(path) else None + + +def read_meta(worktree_path: str) -> Optional[Dict[str, Any]]: + admin = admin_dir(worktree_path) + if not admin: + return None + meta_path = os.path.join(admin, META_NAME) + if not os.path.isfile(meta_path): + return None + try: + with open(meta_path, "r", encoding="utf-8") as fh: + return json.load(fh) + except (ValueError, OSError): + return None + + +def write_meta(worktree_path: str, meta: Dict[str, Any]) -> None: + admin = admin_dir(worktree_path) + if not admin: + return + with open(os.path.join(admin, META_NAME), "w", encoding="utf-8") as fh: + json.dump(meta, fh, indent=2, sort_keys=True) + + +def is_dirty(worktree_path: str) -> bool: + return bool(git(worktree_path, "status", "--porcelain", check=False).strip()) + + +def unpushed(worktree_path: str) -> bool: + """True when HEAD exists on no remote at all. + + Being ahead of the tracked upstream is not enough to call work unpushed. + A review branch commonly tracks the base it will merge into (say + ``upstream/master``) while the commits themselves are pushed to a fork, so + counting ``@{upstream}..HEAD`` alone reports safe work as unsafe. Only if + no remote-tracking ref contains HEAD is anything actually at risk. + """ + ahead = run(["git", "-C", worktree_path, "rev-list", "--count", "@{upstream}..HEAD"], check=False).strip() + try: + if int(ahead) == 0: + return False + except ValueError: + pass # No upstream configured; fall through to the remote check. + contains = run(["git", "-C", worktree_path, "branch", "-r", "--contains", "HEAD"], check=False).strip() + return not contains + + +# --------------------------------------------------------------------------- # +# environment detection +# --------------------------------------------------------------------------- # +def detect_venv(repo: str) -> Optional[str]: + """Return the bin/ directory of the project's virtualenv, if any. + + Resolved against the *main* worktree on purpose: poetry keys its + virtualenvs by project path, so asking from inside a fresh worktree would + point at an env that does not exist. + """ + env_var = os.environ.get("VIRTUAL_ENV") + candidates: List[str] = [] + if shutil.which("poetry"): + out = run(["poetry", "env", "info", "--path"], cwd=repo, check=False).strip() + if out: + candidates.append(out) + name = os.path.basename(repo.rstrip("/")).lower() + candidates += sorted(glob.glob(os.path.join(os.path.expanduser("~"), ".cache", "pypoetry", "virtualenvs", f"{name}-*"))) + candidates += [os.path.join(repo, ".venv")] + if env_var: + candidates.append(env_var) + for cand in candidates: + binary = os.path.join(cand, "bin", "python") + if os.path.isfile(binary): + return os.path.join(cand, "bin") + return None + + +# --------------------------------------------------------------------------- # +# gh integration +# --------------------------------------------------------------------------- # +def pr_head(repo: str, number: int, slug: Optional[str]) -> Dict[str, Any]: + if not shutil.which("gh"): + raise SystemExit("gh CLI not found; use --branch/--from instead of --pr") + cmd = [ + "gh", + "pr", + "view", + str(number), + "--json", + "number,title,state,headRefName,baseRefName,headRepository,headRepositoryOwner,isCrossRepository", + ] + if slug: + cmd += ["--repo", slug] + try: + out = run(cmd, cwd=repo) + except CommandError as exc: + raise SystemExit(f"gh pr view failed: {exc.stderr.strip()}") + return json.loads(out) + + +# --------------------------------------------------------------------------- # +# commands +# --------------------------------------------------------------------------- # +def sanitize(name: str) -> str: + return re.sub(r"[^A-Za-z0-9._-]+", "-", name).strip("-") or "wt" + + +def unique_branch(repo: str, preferred: str) -> str: + taken = checked_out_branches(repo) + candidate = preferred + suffix = 2 + while branch_exists(repo, candidate) or candidate in taken: + candidate = f"{preferred}-{suffix}" + suffix += 1 + return candidate + + +def cmd_new(args: argparse.Namespace) -> Dict[str, Any]: + repo = main_worktree(args.repo) + slug = repo_slug(repo) + warnings: List[str] = [] + + pr_meta: Optional[Dict[str, Any]] = None + start_point: str + upstream: Optional[str] = None + default_name: str + + if args.pr: + pr_meta = pr_head(repo, args.pr, args.slug or slug) + owner = (pr_meta.get("headRepositoryOwner") or {}).get("login") + head_repo = (pr_meta.get("headRepository") or {}).get("name") + head_ref = pr_meta["headRefName"] + default_name = f"pr{args.pr}" + remote = remote_for(repo, owner, head_repo) if owner and head_repo else None + if remote: + git(repo, "fetch", "--quiet", remote, f"{head_ref}:refs/remotes/{remote}/{head_ref}") + start_point = f"{remote}/{head_ref}" + upstream = start_point + else: + url = f"https://github.com/{owner}/{head_repo}.git" + ref = f"{FETCH_NS}/{owner}/{head_ref}" + git(repo, "fetch", "--quiet", url, f"+{head_ref}:{ref}") + start_point = ref + warnings.append( + f"no local remote for {owner}/{head_repo}; fetched from {url}. " + f"No upstream set - `git push` will need an explicit remote." + ) + branch_pref = args.branch_name or head_ref + elif args.branch: + head_ref = args.branch + default_name = sanitize(head_ref) + configured = remotes(repo) + if args.remote: + candidates = [args.remote] + else: + # Try origin first, then the canonical repo, then anything else: + # in a fork workflow the branch may live on either side. + candidates = [name for name in ("origin", "upstream") if name in configured] + candidates += [name for name in configured if name not in candidates] + if not candidates: + raise SystemExit("no remotes configured; pass --remote") + + remote = None + for candidate in candidates: + try: + git(repo, "fetch", "--quiet", candidate, f"{head_ref}:refs/remotes/{candidate}/{head_ref}") + except CommandError: + continue + remote = candidate + break + if not remote: + raise SystemExit(f"branch {head_ref!r} not found on any of: {', '.join(candidates)}") + if remote != candidates[0]: + warnings.append(f"{head_ref} was not on {candidates[0]}; fetched from {remote}") + + start_point = f"{remote}/{head_ref}" + upstream = start_point + branch_pref = args.branch_name or head_ref + elif args.from_ref: + start_point = args.from_ref + default_name = sanitize(args.name or args.from_ref) + branch_pref = args.branch_name or f"agent/{default_name}" + else: + raise SystemExit("one of --pr, --branch or --from is required") + + name = sanitize(args.name or default_name) + base_dir = args.base_dir or DEFAULT_BASE_DIR + path = args.path or os.path.join(base_dir, os.path.basename(repo.rstrip("/")), name) + + if os.path.exists(path): + if not args.force: + raise SystemExit(f"{path} already exists (use --force to replace, or `remove {name}`)") + _drop_worktree(repo, path, force=True, keep_branch=True) + + os.makedirs(os.path.dirname(path), exist_ok=True) + branch = unique_branch(repo, branch_pref) + if branch != branch_pref: + warnings.append(f"branch {branch_pref} was taken; using {branch}") + + git(repo, "worktree", "add", "--quiet", "-b", branch, path, start_point) + if upstream: + git(path, "branch", "--set-upstream-to", upstream, branch, check=False) + + meta = { + "managed_by": "agent-worktree", + "name": name, + "repo": repo, + "branch": branch, + "created_branch": True, + "start_point": start_point, + "upstream": upstream, + "pr": args.pr, + "created_at": _dt.datetime.now().astimezone().isoformat(timespec="seconds"), + } + write_meta(path, meta) + + result: Dict[str, Any] = { + "path": path, + "branch": branch, + "start_point": start_point, + "upstream": upstream, + "head": git(path, "rev-parse", "HEAD").strip(), + "venv_bin": detect_venv(repo), + "warnings": warnings, + } + if pr_meta: + result["pr"] = { + "number": pr_meta["number"], + "title": pr_meta.get("title"), + "state": pr_meta.get("state"), + "head": pr_meta.get("headRefName"), + "base": pr_meta.get("baseRefName"), + "cross_repository": pr_meta.get("isCrossRepository"), + } + return result + + +def _resolve(repo: str, target: str) -> str: + """Map a name or path to an existing worktree path.""" + if os.path.isdir(target): + return os.path.abspath(target) + for entry in worktrees(repo): + meta = read_meta(entry["path"]) + if meta and meta.get("name") == target: + return entry["path"] + if os.path.basename(entry["path"].rstrip("/")) == target: + return entry["path"] + if entry.get("branch") == target: + return entry["path"] + raise SystemExit(f"no worktree matching {target!r}") + + +def _drop_worktree(repo: str, path: str, force: bool, keep_branch: bool) -> Dict[str, Any]: + meta = read_meta(path) or {} + branch = meta.get("branch") + removed_branch = False + cmd = ["worktree", "remove"] + if force: + cmd.append("--force") + cmd.append(path) + try: + git(repo, *cmd) + except CommandError as exc: + if not force: + raise + shutil.rmtree(path, ignore_errors=True) + git(repo, "worktree", "prune", check=False) + _ = exc + git(repo, "worktree", "prune", check=False) + if branch and meta.get("created_branch") and not keep_branch: + rc = subprocess.run( + ["git", "-C", repo, "branch", "-D", branch], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ).returncode + removed_branch = rc == 0 + return {"path": path, "branch": branch, "branch_deleted": removed_branch} + + +def cmd_remove(args: argparse.Namespace) -> Dict[str, Any]: + repo = main_worktree(args.repo) + path = _resolve(repo, args.target) + if path == repo: + raise SystemExit("refusing to remove the main worktree") + meta = read_meta(path) + if not meta and not args.force: + raise SystemExit(f"{path} was not created by agent-worktree (use --force to remove anyway)") + if not args.force: + if is_dirty(path): + raise SystemExit(f"{path} has uncommitted changes (use --force to discard)") + if unpushed(path): + raise SystemExit(f"{path} has unpushed commits (use --force to discard)") + return _drop_worktree(repo, path, force=args.force, keep_branch=args.keep_branch) + + +def cmd_cleanup(args: argparse.Namespace) -> Dict[str, Any]: + repo = main_worktree(args.repo) + cutoff = None + if args.older_than is not None: + cutoff = _dt.datetime.now().astimezone() - _dt.timedelta(days=args.older_than) + + removed: List[Dict[str, Any]] = [] + skipped: List[Dict[str, Any]] = [] + for entry in worktrees(repo): + path = entry["path"] + if path == repo: + continue + meta = read_meta(path) + if not meta: + continue + if cutoff: + try: + created = _dt.datetime.fromisoformat(meta["created_at"]) + except (KeyError, ValueError): + created = None + if created and created > cutoff: + skipped.append({"path": path, "reason": "newer than cutoff"}) + continue + if not args.force: + if is_dirty(path): + skipped.append({"path": path, "reason": "uncommitted changes"}) + continue + if unpushed(path): + skipped.append({"path": path, "reason": "unpushed commits"}) + continue + removed.append(_drop_worktree(repo, path, force=True, keep_branch=args.keep_branch)) + git(repo, "worktree", "prune", check=False) + return {"removed": removed, "skipped": skipped} + + +def cmd_list(args: argparse.Namespace) -> Dict[str, Any]: + repo = main_worktree(args.repo) + rows = [] + for entry in worktrees(repo): + meta = read_meta(entry["path"]) + rows.append( + { + "path": entry["path"], + "branch": entry.get("branch"), + "head": entry.get("head", "")[:12], + "main": entry["path"] == repo, + "managed": bool(meta), + "name": (meta or {}).get("name"), + "pr": (meta or {}).get("pr"), + "created_at": (meta or {}).get("created_at"), + "dirty": is_dirty(entry["path"]), + } + ) + return {"repo": repo, "worktrees": rows} + + +def cmd_path(args: argparse.Namespace) -> Dict[str, Any]: + repo = main_worktree(args.repo) + return {"path": _resolve(repo, args.target)} + + +def cmd_update(args: argparse.Namespace) -> Dict[str, Any]: + repo = main_worktree(args.repo) + path = _resolve(repo, args.target) + meta = read_meta(path) or {} + if is_dirty(path) and not args.force: + raise SystemExit(f"{path} has uncommitted changes (use --force to discard)") + + before = git(path, "rev-parse", "HEAD").strip() + upstream = meta.get("upstream") + start_point = meta.get("start_point") + if upstream and "/" in upstream: + remote, ref = upstream.split("/", 1) + git(repo, "fetch", "--quiet", remote, f"{ref}:refs/remotes/{remote}/{ref}") + target = upstream + elif start_point: + target = start_point + else: + raise SystemExit(f"{path} has no recorded start point; update manually") + + if args.rebase: + git(path, "rebase", target) + else: + git(path, "reset", "--hard", target) + return { + "path": path, + "before": before[:12], + "after": git(path, "rev-parse", "HEAD").strip()[:12], + "target": target, + } + + +def cmd_info(args: argparse.Namespace) -> Dict[str, Any]: + repo = main_worktree(args.repo) + return { + "repo": repo, + "slug": repo_slug(repo), + "git_common_dir": git_common_dir(repo), + "branch": git(repo, "rev-parse", "--abbrev-ref", "HEAD").strip(), + "dirty": is_dirty(repo), + "remotes": remotes(repo), + "venv_bin": detect_venv(repo), + "base_dir": args.base_dir or DEFAULT_BASE_DIR, + "gh": shutil.which("gh"), + } + + +# --------------------------------------------------------------------------- # +# rendering +# --------------------------------------------------------------------------- # +def render(command: str, data: Dict[str, Any]) -> str: + lines: List[str] = [] + if command == "new": + lines.append(f"path {data['path']}") + lines.append(f"branch {data['branch']}") + lines.append(f"from {data['start_point']} @ {data['head'][:12]}") + if data.get("upstream"): + lines.append(f"upstream {data['upstream']}") + if data.get("pr"): + pr = data["pr"] + lines.append(f"pr #{pr['number']} [{pr['state']}] {pr['head']} -> {pr['base']}") + lines.append(f" {pr['title']}") + if data.get("venv_bin"): + lines.append(f"venv {data['venv_bin']}") + for warning in data.get("warnings", []): + lines.append(f"WARNING {warning}") + elif command == "list": + lines.append(f"repo {data['repo']}") + for row in data["worktrees"]: + tags = [] + if row["main"]: + tags.append("main") + if row["managed"]: + tags.append("managed") + if row["dirty"]: + tags.append("dirty") + if row["pr"]: + tags.append(f"pr#{row['pr']}") + suffix = f" [{','.join(tags)}]" if tags else "" + lines.append(f" {row['branch'] or '(detached)':<40} {row['head']} {row['path']}{suffix}") + elif command == "path": + lines.append(data["path"]) + elif command == "remove": + lines.append(f"removed {data['path']}") + if data.get("branch_deleted"): + lines.append(f"deleted branch {data['branch']}") + elif command == "cleanup": + for item in data["removed"]: + lines.append(f"removed {item['path']}") + for item in data["skipped"]: + lines.append(f"kept {item['path']} ({item['reason']})") + if not data["removed"] and not data["skipped"]: + lines.append("nothing to clean") + elif command == "update": + lines.append(f"{data['path']}: {data['before']} -> {data['after']} ({data['target']})") + elif command == "info": + for key in ("repo", "slug", "branch", "dirty", "venv_bin", "base_dir", "gh"): + lines.append(f"{key:<15} {data.get(key)}") + lines.append("remotes") + for name, url in data["remotes"].items(): + lines.append(f" {name:<13} {url}") + else: + lines.append(json.dumps(data, indent=2)) + return "\n".join(lines) + + +# --------------------------------------------------------------------------- # +# cli +# --------------------------------------------------------------------------- # +def build_parser() -> argparse.ArgumentParser: + # Global flags are attached to every subparser as well, so they work both + # before and after the subcommand. argparse.SUPPRESS keeps an unspecified + # subparser flag from clobbering a value given before the subcommand. + common = argparse.ArgumentParser(add_help=False) + common.add_argument("--repo", default=argparse.SUPPRESS, help="path inside the target repository (default: cwd)") + common.add_argument( + "--base-dir", + default=argparse.SUPPRESS, + help=f"where worktrees live (default: {DEFAULT_BASE_DIR})", + ) + common.add_argument("--json", action="store_true", default=argparse.SUPPRESS, help="emit JSON") + + parser = argparse.ArgumentParser( + prog="agent_worktree.py", + description="Create and manage disposable git worktrees for agent work.", + parents=[common], + ) + # No set_defaults() for the shared flags: argparse's `parents` mechanism + # shares the *same* action objects with every subparser, so setting a + # default here would also reset the subparser copy and silently discard a + # flag passed before the subcommand. Defaults are applied in main(). + sub = parser.add_subparsers(dest="command", required=True) + + def add(name: str, **kwargs: Any) -> argparse.ArgumentParser: + return sub.add_parser(name, parents=[common], **kwargs) + + p_new = add("new", aliases=["create"], help="create a worktree") + source = p_new.add_mutually_exclusive_group(required=True) + source.add_argument("--pr", type=int, help="pull request number (resolved via gh)") + source.add_argument("--branch", help="existing remote branch name") + source.add_argument("--from", dest="from_ref", help="any git ref to branch from") + p_new.add_argument("--name", help="short worktree name (default: pr or branch name)") + p_new.add_argument("--branch-name", help="local branch name to create") + p_new.add_argument("--remote", help="remote to fetch from (with --branch)") + p_new.add_argument("--slug", help="owner/repo override for gh") + p_new.add_argument("--path", help="explicit worktree path") + p_new.add_argument("--force", action="store_true", help="replace an existing path") + p_new.set_defaults(func=cmd_new) + + p_list = add("list", aliases=["ls"], help="list worktrees") + p_list.set_defaults(func=cmd_list) + + p_path = add("path", help="print the path of a worktree") + p_path.add_argument("target") + p_path.set_defaults(func=cmd_path) + + p_update = add("update", aliases=["sync"], help="fast-forward a worktree to its source") + p_update.add_argument("target") + p_update.add_argument("--rebase", action="store_true", help="rebase instead of reset --hard") + p_update.add_argument("--force", action="store_true", help="discard local changes") + p_update.set_defaults(func=cmd_update) + + p_remove = add("remove", aliases=["rm"], help="remove a worktree") + p_remove.add_argument("target") + p_remove.add_argument("--force", action="store_true", help="discard uncommitted/unpushed work") + p_remove.add_argument("--keep-branch", action="store_true", help="do not delete the local branch") + p_remove.set_defaults(func=cmd_remove) + + p_clean = add("cleanup", help="remove all managed worktrees") + p_clean.add_argument("--all", action="store_true", help="no-op, kept for readability") + p_clean.add_argument("--older-than", type=int, metavar="DAYS") + p_clean.add_argument("--force", action="store_true", help="discard uncommitted/unpushed work") + p_clean.add_argument("--keep-branch", action="store_true") + p_clean.set_defaults(func=cmd_cleanup) + + p_info = add("info", help="show repository, remotes and venv") + p_info.set_defaults(func=cmd_info) + + return parser + + +def main(argv: Optional[List[str]] = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + # The shared flags use argparse.SUPPRESS so that a value given before the + # subcommand survives; backfill whatever was never supplied. + for name, default in (("repo", None), ("base_dir", None), ("json", False)): + if not hasattr(args, name): + setattr(args, name, default) + try: + data = args.func(args) + except CommandError as exc: + print(str(exc), file=sys.stderr) + return 1 + except SystemExit as exc: + if isinstance(exc.code, str): + print(exc.code, file=sys.stderr) + return 1 + raise + canonical = {"create": "new", "ls": "list", "rm": "remove", "sync": "update"} + command = canonical.get(args.command, args.command) + if args.json: + print(json.dumps(data, indent=2, sort_keys=True)) + else: + print(render(command, data)) + return 0 + + +if __name__ == "__main__": + sys.exit(main())