gitmaster.py

#!/usr/bin/env python3
"""
GitMaster

Recover source files from an exposed .git/ directory, including repositories
that store objects as loose objects, packfiles, or a mix of both.
"""

from __future__ import annotations

import argparse
import concurrent.futures
import sys
import re
import shutil
import ssl
import struct
import subprocess
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import zlib

from collections import Counter
from dataclasses import dataclass
from pathlib import Path, PurePosixPath
from typing import Optional

try:
    from tqdm import tqdm
except Exception:  # pragma: no cover - optional dependency fallback
    tqdm = None


USER_AGENT = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
    "AppleWebKit/537.36 (KHTML, like Gecko) "
    "Chrome/124.0 Safari/537.36"
)

PACK_NAME_RE = re.compile(r"pack-[0-9a-fA-F]{40}\.pack")


class Ansi:
    RESET = "\033[0m"
    BOLD = "\033[1m"
    DIM = "\033[2m"
    RED = "\033[31m"
    GREEN = "\033[32m"
    YELLOW = "\033[33m"
    BLUE = "\033[34m"
    MAGENTA = "\033[35m"
    CYAN = "\033[36m"


@dataclass(frozen=True)
class IndexEntry:
    path: str
    sha1: str
    mode: int
    size: int


@dataclass(frozen=True)
class RecoveryResult:
    status: str
    path: str
    detail: str = ""


class Log:
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self.color = self._supports_color()

    def write(self, message: str) -> None:
        with self._lock:
            print(message, flush=True)

    def write_status(self, label: str, message: str, color: str = "") -> None:
        with self._lock:
            prefix = self._format_label(label, color)
            print(f"{prefix} {message}", flush=True)

    def colorize(self, text: str, color: str = "", bold: bool = False) -> str:
        if not self.color or not color:
            return text
        style = color
        if bold:
            style = Ansi.BOLD + color
        return f"{style}{text}{Ansi.RESET}"

    def success_mark(self) -> str:
        return self.colorize("✓", Ansi.GREEN, bold=True)

    def failure_mark(self) -> str:
        return self.colorize("✗", Ansi.RED, bold=True)

    def _format_label(self, label: str, color: str) -> str:
        if self.color and color:
            return f"{Ansi.BOLD}{color}[{label}]{Ansi.RESET}"
        return f"[{label}]"

    @staticmethod
    def _supports_color() -> bool:
        return hasattr(sys.stdout, "isatty") and sys.stdout.isatty()


class ProgressDisplay:
    STATUS_COLORS = {
        "OK": Ansi.GREEN,
        "MISSING": Ansi.YELLOW,
        "ERROR": Ansi.RED,
    }

    def __init__(self, total: int, color_enabled: bool) -> None:
        self.total = total
        self.count = 0
        self._lock = threading.Lock()
        self._bar = None
        self.color_enabled = color_enabled
        if tqdm is not None:
            self._bar = tqdm(
                total=total,
                desc="Recovering",
                unit="file",
                dynamic_ncols=True,
                colour="green",
                leave=True,
                file=sys.stdout,
            )

    def update(self, result: RecoveryResult) -> None:
        with self._lock:
            self.count += 1
            if self._bar is not None:
                status_text = self._style_status(result.status)
                detail = f"{status_text}: {self._shorten(result.path)}"
                self._bar.set_postfix_str(detail, refresh=False)
                self._bar.update(1)
            else:
                print(
                    f"[{self.count}/{self.total}] {self._style_status(result.status)}: {result.path}",
                    flush=True,
                )

    def close(self) -> None:
        with self._lock:
            if self._bar is not None:
                self._bar.close()

    @staticmethod
    def _shorten(path: str, max_length: int = 48) -> str:
        if len(path) <= max_length:
            return path
        return "..." + path[-(max_length - 3):]

    def _style_status(self, status: str) -> str:
        if not self.color_enabled:
            return status
        color = self.STATUS_COLORS.get(status, "")
        if not color:
            return status
        return f"{Ansi.BOLD}{color}{status}{Ansi.RESET}"


class SpinnerTask:
    FRAMES = ["|", "/", "-", "\\"]

    def __init__(self, logger: Log, label: str, message: str, color: str = "") -> None:
        self.logger = logger
        self.label = label
        self.message = message
        self.color = color
        self.enabled = logger.color
        self._stop = threading.Event()
        self._thread: Optional[threading.Thread] = None
        self._start_time = 0.0
        self._finished = False

    def __enter__(self) -> "SpinnerTask":
        self.start()
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        if self._finished:
            return
        if exc is None:
            self.success()
        else:
            self.fail(str(exc))

    def start(self) -> None:
        self._start_time = time.time()
        if not self.enabled:
            self.logger.write_status(self.label, f"{self.message} ...", self.color)
            return
        self._thread = threading.Thread(target=self._spin, daemon=True)
        self._thread.start()

    def success(self, detail: Optional[str] = None) -> None:
        self._finish(True, detail)

    def fail(self, detail: Optional[str] = None) -> None:
        self._finish(False, detail)

    def _spin(self) -> None:
        index = 0
        while not self._stop.wait(0.12):
            with self.logger._lock:
                prefix = self.logger._format_label(self.label, self.color)
                frame = self.logger.colorize(self.FRAMES[index], Ansi.CYAN, bold=True)
                sys.stdout.write("\r\033[K" + f"{prefix} {self.message} {frame}")
                sys.stdout.flush()
            index = (index + 1) % len(self.FRAMES)

    def _finish(self, success: bool, detail: Optional[str]) -> None:
        if self._finished:
            return
        self._finished = True
        elapsed = time.time() - self._start_time if self._start_time else 0.0
        message = detail or self.message
        mark = self.logger.success_mark() if success else self.logger.failure_mark()
        if self.enabled:
            self._stop.set()
            if self._thread is not None:
                self._thread.join()
            with self.logger._lock:
                prefix = self.logger._format_label(self.label, self.color)
                sys.stdout.write(
                    "\r\033[K" + f"{prefix} {message} {mark} {self.logger.colorize(f'({elapsed:.1f}s)', Ansi.DIM)}\n"
                )
                sys.stdout.flush()
        else:
            self.logger.write_status(self.label, f"{message} {mark} ({elapsed:.1f}s)", self.color)


class HttpClient:
    def __init__(self, base_url: str, timeout: float = 10.0) -> None:
        self.base_url = base_url.rstrip("/") + "/"
        self.timeout = timeout
        self.context = ssl._create_unverified_context()

    def _build_url(self, relative_path: str) -> str:
        return urllib.parse.urljoin(self.base_url, relative_path.lstrip("/"))

    def fetch(self, relative_path: str) -> bytes:
        url = self._build_url(relative_path)
        request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
        with urllib.request.urlopen(request, timeout=self.timeout, context=self.context) as response:
            return response.read()

    def fetch_optional(self, relative_path: str) -> Optional[bytes]:
        try:
            return self.fetch(relative_path)
        except urllib.error.HTTPError as exc:
            if exc.code == 404:
                return None
            raise


class GitIndexParser:
    def __init__(self, data: bytes) -> None:
        self.data = data
        self.offset = 0

    def read(self, size: int) -> bytes:
        if self.offset + size > len(self.data):
            raise ValueError("unexpected end of index data")
        chunk = self.data[self.offset:self.offset + size]
        self.offset += size
        return chunk

    def read_u32(self) -> int:
        return struct.unpack("!I", self.read(4))[0]

    def read_u16(self) -> int:
        return struct.unpack("!H", self.read(2))[0]

    def parse(self) -> list[IndexEntry]:
        signature = self.read(4)
        if signature != b"DIRC":
            raise ValueError("not a Git index file")

        version = self.read_u32()
        if version not in (2, 3):
            raise ValueError(f"unsupported Git index version: {version}")

        entry_count = self.read_u32()
        entries: list[IndexEntry] = []

        for _ in range(entry_count):
            entry_start = self.offset

            # stat data
            self.read(4 * 6)
            mode = self.read_u32()
            self.read(4 * 2)
            size = self.read_u32()
            sha1 = self.read(20).hex()
            flags = self.read_u16()

            extended = bool(flags & 0x4000)
            name_length = flags & 0x0FFF

            if extended and version == 3:
                self.read_u16()

            if name_length < 0x0FFF:
                raw_name = self.read(name_length)
            else:
                name_bytes = bytearray()
                while True:
                    ch = self.read(1)
                    if ch == b"\x00":
                        break
                    name_bytes.extend(ch)
                raw_name = bytes(name_bytes)

            consumed = self.offset - entry_start
            padding_length = (8 - (consumed % 8)) or 8
            padding = self.read(padding_length)
            if any(byte != 0 for byte in padding):
                raise ValueError("index padding contained non-NUL bytes")

            path = raw_name.decode("utf-8", errors="replace")
            entries.append(IndexEntry(path=path, sha1=sha1, mode=mode, size=size))

        return entries


class LocalGitMirror:
    def __init__(self, workspace: Path, logger: Log) -> None:
        self.workspace = workspace
        self.git_dir = workspace / ".git"
        self.logger = logger
        self.git_binary = shutil.which("git")
        self.git_ready = False

    def initialize(self) -> None:
        self.workspace.mkdir(parents=True, exist_ok=True)
        if self.git_binary:
            try:
                subprocess.run(
                    [self.git_binary, "init", "-q", str(self.workspace)],
                    check=True,
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                )
                self.git_ready = True
            except subprocess.SubprocessError:
                self.git_ready = False

        for relative in (
            Path("."),
            Path("objects"),
            Path("objects/info"),
            Path("objects/pack"),
            Path("refs"),
        ):
            (self.git_dir / relative).mkdir(parents=True, exist_ok=True)

    def write_git_file(self, relative_path: str, content: bytes) -> None:
        target = self.git_dir / Path(relative_path)
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_bytes(content)

    def store_loose_object(self, sha1: str, compressed_data: bytes) -> None:
        target = self.git_dir / "objects" / sha1[:2] / sha1[2:]
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_bytes(compressed_data)

    def has_object(self, sha1: str) -> bool:
        if self.git_ready:
            result = subprocess.run(
                [self.git_binary, f"--git-dir={self.git_dir}", "cat-file", "-e", sha1],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            return result.returncode == 0
        return (self.git_dir / "objects" / sha1[:2] / sha1[2:]).exists()

    def read_blob(self, sha1: str) -> bytes:
        if self.git_ready:
            result = subprocess.run(
                [self.git_binary, f"--git-dir={self.git_dir}", "cat-file", "-p", sha1],
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
            )
            return result.stdout

        object_path = self.git_dir / "objects" / sha1[:2] / sha1[2:]
        payload = zlib.decompress(object_path.read_bytes())
        _header, content = payload.split(b"\x00", 1)
        return content

    def index_pack(self, pack_path: Path) -> bool:
        if not self.git_ready:
            return False
        try:
            subprocess.run(
                [self.git_binary, f"--git-dir={self.git_dir}", "index-pack", str(pack_path)],
                check=True,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            return True
        except subprocess.SubprocessError:
            return False


class Recoverer:
    def __init__(
        self,
        target_url: str,
        output_dir: Optional[str],
        threads: int,
        timeout: float,
    ) -> None:
        self.target_url = target_url.rstrip("/") + "/"
        self.client = HttpClient(self.target_url, timeout=timeout)
        self.logger = Log()
        self.threads = max(1, threads)
        self.stats = Counter()
        self.stats_lock = threading.Lock()

        parsed = urllib.parse.urlparse(self.target_url)
        default_name = parsed.netloc.replace(":", "_") or "recovered_repo"
        self.workspace = Path(output_dir or default_name).resolve()
        self.mirror = LocalGitMirror(self.workspace, self.logger)
        self.packfiles: list[str] = []
        self.results: dict[str, RecoveryResult] = {}

    def run(self) -> int:
        with SpinnerTask(self.logger, "INIT", f"Preparing workspace {self.workspace.name}", Ansi.BLUE):
            self.mirror.initialize()
        self._mirror_metadata()
        entries = self._load_index_entries()
        if not entries:
            self.logger.write("[!] No valid entries found in .git/index")
            return 1

        with concurrent.futures.ThreadPoolExecutor(max_workers=self.threads) as executor:
            progress = ProgressDisplay(total=len(entries), color_enabled=self.logger.color)
            try:
                futures = [executor.submit(self._recover_entry, entry) for entry in entries]
                for future in concurrent.futures.as_completed(futures):
                    result = future.result()
                    self.results[result.path] = result
                    progress.update(result)
            finally:
                progress.close()

        self._print_source_tree(entries)
        self.logger.write_status(
            "DONE",
            f"recovered={self.stats['recovered']} missing={self.stats['missing']} "
            f"errors={self.stats['errors']} output={self.workspace}",
            Ansi.CYAN,
        )
        return 0 if self.stats["recovered"] else 1

    def _mirror_metadata(self) -> None:
        with SpinnerTask(self.logger, "META", "Fetching HEAD", Ansi.BLUE) as spinner:
            head = self.client.fetch_optional("HEAD")
            spinner.success("Fetched HEAD" if head else "HEAD not exposed")
        if head:
            self.mirror.write_git_file("HEAD", head)
            ref_name = self._parse_head_ref(head)
            if ref_name:
                with SpinnerTask(self.logger, "META", f"Fetching {ref_name}", Ansi.BLUE) as spinner:
                    ref_content = self.client.fetch_optional(ref_name)
                    spinner.success(f"Fetched {ref_name}" if ref_content else f"{ref_name} not exposed")
                if ref_content:
                    self.mirror.write_git_file(ref_name, ref_content)

        with SpinnerTask(self.logger, "META", "Fetching packed-refs", Ansi.BLUE) as spinner:
            packed_refs = self.client.fetch_optional("packed-refs")
            spinner.success("Fetched packed-refs" if packed_refs else "packed-refs not exposed")
        if packed_refs:
            self.mirror.write_git_file("packed-refs", packed_refs)

        with SpinnerTask(self.logger, "META", "Fetching objects/info/packs", Ansi.BLUE) as spinner:
            pack_listing = self.client.fetch_optional("objects/info/packs")
            spinner.success("Fetched objects/info/packs" if pack_listing else "No pack listing exposed")
        if pack_listing:
            self.mirror.write_git_file("objects/info/packs", pack_listing)
            self.packfiles = self._parse_pack_listing(pack_listing)
            for pack_name in self.packfiles:
                self._mirror_pack(pack_name)

    def _mirror_pack(self, pack_name: str) -> None:
        with SpinnerTask(self.logger, "PACK", f"Downloading {pack_name}", Ansi.MAGENTA) as spinner:
            pack_bytes = self.client.fetch_optional(f"objects/pack/{pack_name}")
            spinner.success(f"Downloaded {pack_name}" if pack_bytes else f"{pack_name} not exposed")
        if not pack_bytes:
            return

        pack_rel = f"objects/pack/{pack_name}"
        self.mirror.write_git_file(pack_rel, pack_bytes)

        idx_name = pack_name[:-5] + ".idx"
        with SpinnerTask(self.logger, "PACK", f"Downloading {idx_name}", Ansi.MAGENTA) as spinner:
            idx_bytes = self.client.fetch_optional(f"objects/pack/{idx_name}")
            spinner.success(f"Downloaded {idx_name}" if idx_bytes else f"{idx_name} not exposed, trying index-pack")
        if idx_bytes:
            self.mirror.write_git_file(f"objects/pack/{idx_name}", idx_bytes)
        else:
            pack_path = self.mirror.git_dir / "objects/pack" / pack_name
            with SpinnerTask(self.logger, "PACK", f"Indexing {pack_name}", Ansi.MAGENTA) as spinner:
                indexed = self.mirror.index_pack(pack_path)
                if indexed:
                    spinner.success(f"Indexed {pack_name}")
                else:
                    spinner.fail(f"Failed to index {pack_name}")
            if not indexed:
                self.logger.write_status("WARN", f"failed to index pack {pack_name}", Ansi.YELLOW)

    def _load_index_entries(self) -> list[IndexEntry]:
        with SpinnerTask(self.logger, "INFO", "Downloading and parsing .git/index", Ansi.BLUE) as spinner:
            data = self.client.fetch("index")
            spinner.success("Downloaded and parsed .git/index")
        self.mirror.write_git_file("index", data)
        parser = GitIndexParser(data)

        entries: list[IndexEntry] = []
        for entry in parser.parse():
            if self._is_safe_path(entry.path):
                entries.append(entry)
            else:
                self.logger.write_status("WARN", f"skipped suspicious path: {entry.path}", Ansi.YELLOW)
        self.logger.write_status("INFO", f"Indexed {len(entries)} file(s)", Ansi.BLUE)
        return entries

    def _recover_entry(self, entry: IndexEntry) -> RecoveryResult:
        try:
            if self.mirror.has_object(entry.sha1):
                blob = self.mirror.read_blob(entry.sha1)
                self._write_output_file(entry.path, blob)
                self._count("recovered")
                return RecoveryResult(status="OK", path=entry.path)

            loose = self.client.fetch_optional(f"objects/{entry.sha1[:2]}/{entry.sha1[2:]}")
            if loose is not None:
                self.mirror.store_loose_object(entry.sha1, loose)
                blob = self.mirror.read_blob(entry.sha1)
                self._write_output_file(entry.path, blob)
                self._count("recovered")
                return RecoveryResult(status="OK", path=entry.path)

            if self.packfiles and self.mirror.has_object(entry.sha1):
                blob = self.mirror.read_blob(entry.sha1)
                self._write_output_file(entry.path, blob)
                self._count("recovered")
                return RecoveryResult(status="OK", path=entry.path)

            self._count("missing")
            return RecoveryResult(status="MISSING", path=entry.path)
        except Exception as exc:
            self._count("errors")
            return RecoveryResult(status="ERROR", path=entry.path, detail=str(exc))

    def _write_output_file(self, relative_path: str, content: bytes) -> None:
        safe_target = self.workspace / self._safe_path(relative_path)
        safe_target.parent.mkdir(parents=True, exist_ok=True)
        safe_target.write_bytes(content)

    def _count(self, key: str) -> None:
        with self.stats_lock:
            self.stats[key] += 1

    @staticmethod
    def _parse_head_ref(head_bytes: bytes) -> Optional[str]:
        try:
            text = head_bytes.decode("utf-8", errors="replace").strip()
        except Exception:
            return None
        if text.startswith("ref: "):
            return text[5:].strip()
        return None

    @staticmethod
    def _parse_pack_listing(data: bytes) -> list[str]:
        try:
            text = data.decode("utf-8", errors="replace")
        except Exception:
            return []
        found = PACK_NAME_RE.findall(text)
        # keep order but remove duplicates
        unique: list[str] = []
        seen = set()
        for name in found:
            if name not in seen:
                seen.add(name)
                unique.append(name)
        return unique

    @staticmethod
    def _safe_path(path_text: str) -> Path:
        pure = PurePosixPath(path_text)
        if pure.is_absolute():
            raise ValueError("absolute paths are not allowed")
        parts = [part for part in pure.parts if part not in ("", ".")]
        if any(part == ".." for part in parts):
            raise ValueError("path traversal is not allowed")
        if not parts:
            raise ValueError("empty path")
        return Path(*parts)

    def _is_safe_path(self, path_text: str) -> bool:
        try:
            self._safe_path(path_text)
            return True
        except ValueError:
            return False

    def _print_source_tree(self, entries: list[IndexEntry]) -> None:
        tree: dict[str, object] = {}
        for entry in sorted(entries, key=lambda item: item.path.lower()):
            node = tree
            parts = PurePosixPath(entry.path).parts
            for part in parts[:-1]:
                node = node.setdefault(part, {})  # type: ignore[assignment]
            node[parts[-1]] = self.results.get(entry.path, RecoveryResult("MISSING", entry.path))  # type: ignore[index]

        self.logger.write_status("TREE", "Source tree from .git/index", Ansi.CYAN)
        self.logger.write(self.logger.colorize("source/", Ansi.CYAN, bold=True))
        self._render_tree(tree, "")

    def _render_tree(self, node: dict[str, object], prefix: str) -> None:
        dirs: list[tuple[str, dict[str, object]]] = []
        files: list[tuple[str, RecoveryResult]] = []
        for name, value in node.items():
            if isinstance(value, dict):
                dirs.append((name, value))
            else:
                files.append((name, value))

        items = sorted(dirs, key=lambda item: item[0].lower()) + sorted(files, key=lambda item: item[0].lower())
        for index, (name, value) in enumerate(items):
            is_last = index == len(items) - 1
            branch = "└── " if is_last else "├── "
            next_prefix = prefix + ("    " if is_last else "│   ")
            if isinstance(value, dict):
                self.logger.write(f"{prefix}{branch}{name}/")
                self._render_tree(value, next_prefix)
            else:
                mark = self.logger.success_mark() if value.status == "OK" else self.logger.failure_mark()
                self.logger.write(f"{prefix}{branch}{name} {mark}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Recover files from an exposed .git/ directory."
    )
    parser.add_argument("url", help="Base URL pointing to the exposed .git/ directory")
    parser.add_argument(
        "-o",
        "--output",
        help="Output directory for recovered files (default: derived from host)",
    )
    parser.add_argument(
        "-t",
        "--threads",
        type=int,
        default=8,
        help="Number of worker threads (default: 8)",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=10.0,
        help="HTTP timeout in seconds (default: 10)",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    recoverer = Recoverer(
        target_url=args.url,
        output_dir=args.output,
        threads=args.threads,
        timeout=args.timeout,
    )
    try:
        return recoverer.run()
    except KeyboardInterrupt:
        print("\n[!] Interrupted by user", flush=True)
        return 130


if __name__ == "__main__":
    raise SystemExit(main())