From 24c91a589df7e83ebb932916b075c1e1ae3d685f Mon Sep 17 00:00:00 2001 From: Karl Hiramoto Date: Mon, 21 Sep 2026 08:16:19 +0200 Subject: [PATCH 1/2] Fix ResultServer file collisions with SHA256 deduplication and versioning (#3245) * Fix ResultServer file collisions with SHA256 deduplication and versioning * Fix CAPE processing crash when [magika] config section is missing and preserve dropfile name order * Use fast mocks in test_cape_processing_without_magika_config --- .../modules/auxiliary/watchdownloads.py | 29 +- .../common/integrations/file_extra_info.py | 3 +- lib/cuckoo/core/resultserver.py | 302 +++++++++++--- modules/processing/CAPE.py | 42 +- tests/test_resultserver_dedup.py | 369 ++++++++++++++++++ 5 files changed, 669 insertions(+), 76 deletions(-) create mode 100644 tests/test_resultserver_dedup.py 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/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_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", + ] + + From e456bc085140c27bc2f0a2d942d8cc9a7171b359 Mon Sep 17 00:00:00 2001 From: doomedraven Date: Mon, 21 Sep 2026 08:19:17 +0200 Subject: [PATCH 2/2] utils: add agent_worktree.py for isolated PR/branch checkouts (#3228) * Add utils/agent_worktree.py for isolated PR/branch checkouts Reviewing a PR or reproducing a bug against another branch currently means checking it out in your working clone. Most clones carry uncommitted work, so that is either unsafe or requires manual git worktree plumbing: figuring out which fork the head branch lives in, fetching it under a sane name, setting upstream, and cleaning up afterwards. utils/agent_worktree.py wraps that: new --pr resolve the head fork via gh, fetch, branch, check out new --branch / --from list / path / update / remove / cleanup / info It is stdlib-only and repository-agnostic. Worktrees it creates are tagged with metadata inside the git admin directory, so nothing extra appears in git status and cleanup only ever touches its own worktrees. remove and cleanup refuse to discard uncommitted changes or unpushed commits without --force, and the main worktree can never be removed. tests/test_agent_worktree.py covers the CLI against throwaway local repositories; no network, credentials or CAPE configuration required. SKILLS.md gains sections on the local development environment, isolated checkouts, and codebase behaviours that cause silent failures. * agent_worktree: handle fork remote layouts Clones that follow the fork convention have origin pointing at the contributor's fork and upstream at the canonical repository. Two things broke there: * repo_slug() preferred origin, so 'gh pr view' was asked about the fork instead of the repository the PR was opened against. upstream now wins when it is configured; clones made straight from the canonical repo are unaffected because they have no upstream. * 'new --branch' fetched from origin only and failed outright if the branch lived upstream. It now tries origin, then upstream, then any other remote, and reports which one it used. A missing branch is now a clear error naming the remotes that were tried. * agent_worktree: do not call work unpushed when it lives on another remote The removal guard compared HEAD against the tracked upstream only. A review branch normally tracks the branch it will merge into while its commits are pushed to a fork, so 'rev-list @{upstream}..HEAD' is non-zero even though nothing is at risk, and cleanup refused to remove a worktree whose work was safely published. The guard now treats work as unpushed only when no remote-tracking ref contains HEAD at all, which is the condition that actually matters. --- SKILLS.md | 100 +++++ changelog.md | 1 + tests/test_agent_worktree.py | 317 ++++++++++++++ utils/agent_worktree.py | 789 +++++++++++++++++++++++++++++++++++ 4 files changed, 1207 insertions(+) create mode 100644 tests/test_agent_worktree.py create mode 100755 utils/agent_worktree.py 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/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/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/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())