vextoly/gitman

"""Repository storage admission checks shared by imports, forks and Git hooks."""
import os
from pathlib import Path
import stat
import sys

RESERVE_BYTES = 64 * 1024


def repository_bytes(path):
    total = 0
    def failed_scan(error):
        raise error
    for directory, dirs, files in os.walk(path, followlinks=False, onerror=failed_scan):
        for name in files:
            try:
                info = (Path(directory) / name).lstat()
            except FileNotFoundError:
                continue  # A concurrent receive may have removed its quarantine.
            if stat.S_ISREG(info.st_mode):
                total += info.st_size
    return total


def enforce_repository_size(path, limit, reserve=RESERVE_BYTES):
    size = repository_bytes(path)
    if size + reserve > limit:
        raise ValueError(f"Repository size limit is {limit / 1048576:g} MiB. "
                         "This operation would exceed the limit (including room for Git metadata).")


def install_quota_hook(path, limit):
    hook = Path(path) / 'hooks' / 'pre-receive'
    content = ('#!/usr/bin/env python3\n# GitMan repository storage quota\n'
               f'import sys\nsys.path.insert(0, {str(Path(__file__).resolve().parent)!r})\n'
               f'from repo_quota import check_receive\nsys.exit(check_receive({int(limit)}))\n')
    if hook.exists():
        previous = hook.read_text()
        if previous == content and os.access(hook, os.X_OK):
            return
        if 'GitMan repository storage quota' not in previous:
            raise ValueError('Existing pre-receive hook requires manual integration; it was not overwritten.')
    # Use a unique temporary file so concurrent receive setup cannot truncate it.
    import tempfile
    fd, temporary = tempfile.mkstemp(prefix='.quota-', dir=hook.parent)
    try:
        with os.fdopen(fd, 'w') as stream:
            stream.write(content)
        os.chmod(temporary, 0o755)
        os.replace(temporary, hook)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def check_receive(limit):
    updates = [line.split() for line in sys.stdin if line.strip()]
    if updates and all(len(row) == 3 and set(row[1]) == {'0'} for row in updates):
        return 0  # Deleting references remains possible even at the limit.
    try:
        # receive-pack quarantine resides inside objects/, so this counts both
        # existing storage and new objects before any branch/tag is accepted.
        reserve = RESERVE_BYTES + sum(len(row[-1].encode()) * 2 + 1024 for row in updates)
        enforce_repository_size(Path.cwd(), limit, reserve)
    except (OSError, ValueError) as exc:
        print(f'GitMan: {exc}', file=sys.stderr)
        return 1
    return 0