apt-sync.py

#!/usr/bin/env python3
import argparse
import bz2
import concurrent.futures
import gzip
import hashlib
import logging
import lzma
import os
import re
import shutil
import socket
import sys
import threading
import time
import traceback
from email.utils import parsedate_to_datetime
from pathlib import Path
from typing import Dict, List, Tuple

import requests

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
formatter = logging.Formatter(
    "%(asctime)s.%(msecs)03d - %(filename)s:%(lineno)d [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%dT%H:%M:%S",
)
handler.setFormatter(formatter)
logger.addHandler(handler)

APT_SYNC_USER_AGENT = os.getenv("APT_SYNC_USER_AGENT", "APT-Mirror-Tool/1.0")
requests.utils.default_user_agent = lambda: APT_SYNC_USER_AGENT

SESSION_PROXY = ""
thread_local = threading.local()

# set preferred address family
import requests.packages.urllib3.util.connection as urllib3_cn

USE_ADDR_FAMILY = os.getenv("USE_ADDR_FAMILY", "").strip().lower()
if USE_ADDR_FAMILY != "":
    assert USE_ADDR_FAMILY in [
        "ipv4",
        "ipv6",
    ], "USE_ADDR_FAMILY must be either ipv4 or ipv6"
    urllib3_cn.allowed_gai_family = lambda: (
        socket.AF_INET if USE_ADDR_FAMILY == "ipv4" else socket.AF_INET6
    )

OS_TEMPLATE = {
    "ubuntu-lts": ["jammy", "noble"],
    "ubuntu-lts-all": ["xenial", "bionic", "focal", "jammy", "noble"],
    "ubuntu-all": ["focal", "jammy", "lunar", "mantic", "noble", "oracular", "plucky", "questing"],
    "debian-current": ["bullseye", "bookworm", "trixie"],
    "debian-all": ["jessie", "stretch", "buster", "bullseye", "bookworm", "trixie"],
    "debian-latest2": ["bookworm", "trixie"],
    "debian-latest": ["trixie"],
}
ARCH_TEMPLATE = {
    "all": ["amd64", "i386", "arm64", "armhf", "armel", "ppc64el", "s390x", "riscv64"],
    "common": ["amd64", "i386", "arm64", "armhf"],
    "x86": ["amd64", "i386"],
}
ARCH_NO_PKGIDX = ["dep11", "i18n", "cnf", "neon"]
MAX_RETRY = int(os.getenv("MAX_RETRY", "3"))
DOWNLOAD_TIMEOUT = int(os.getenv("DOWNLOAD_TIMEOUT", "1800"))
RETRY_WAIT_SECONDS = float(os.getenv("RETRY_WAIT_SECONDS", "2"))
DOWNLOAD_WORKERS = max(1, int(os.getenv("DOWNLOAD_WORKERS", "4")))
REPO_SIZE_FILE = os.getenv("REPO_SIZE_FILE", "")

pattern_os_template = re.compile(r"@\{(.+)\}")
pattern_package_name = re.compile(r"^Filename: (.+)$", re.MULTILINE)
pattern_package_size = re.compile(r"^Size: (\d+)$", re.MULTILINE)
pattern_package_sha256 = re.compile(r"^SHA256: (\w{64})$", re.MULTILINE)
download_cache = dict()
progress_lock = threading.Lock()


def new_session() -> requests.Session:
    s = requests.Session()
    s.headers.update({"User-Agent": APT_SYNC_USER_AGENT})
    if SESSION_PROXY:
        s.proxies = {"http": SESSION_PROXY, "https": SESSION_PROXY}
    return s


def get_session() -> requests.Session:
    if not hasattr(thread_local, "session"):
        thread_local.session = new_session()
    return thread_local.session


def is_interactive_shell() -> bool:
    return sys.stderr.isatty()


def format_size(size: int) -> str:
    units = ["B", "KiB", "MiB", "GiB", "TiB"]
    val = float(size)
    for unit in units:
        if val < 1024.0 or unit == units[-1]:
            if unit == "B":
                return f"{int(val)}{unit}"
            return f"{val:.1f}{unit}"
        val /= 1024.0
    return f"{size}B"


def print_progress(
    url: str,
    written: int,
    total: int,
    start_ts: float,
    status: str,
    attempt: int,
    max_attempt: int,
    done: bool = False,
):
    if not is_interactive_shell():
        return
    filename = url.rsplit("/", 1)[-1]
    elapsed = max(time.time() - start_ts, 0.001)
    speed = written / elapsed
    if total > 0:
        pct = min(100.0, (written / total) * 100)
        msg = (
            f"downloading {filename} [{attempt}/{max_attempt}] "
            f"{pct:6.2f}% ({format_size(written)}/{format_size(total)}) "
            f"{format_size(int(speed))}/s {status}"
        )
    else:
        msg = (
            f"downloading {filename} [{attempt}/{max_attempt}] "
            f"{format_size(written)} {format_size(int(speed))}/s {status}"
        )
    with progress_lock:
        if done:
            sys.stderr.write(f"\r{msg}\n")
        else:
            sys.stderr.write(f"\r{msg}")
        sys.stderr.flush()


def check_args(prop: str, lst: List[str]):
    for s in lst:
        if len(s) == 0 or " " in s:
            raise ValueError(f"Invalid item in {prop}: {repr(s)}")


def replace_os_template(os_list: List[str]) -> List[str]:
    ret = []
    for i in os_list:
        matched = pattern_os_template.search(i)
        if matched:
            for os in OS_TEMPLATE[matched.group(1)]:
                ret.append(pattern_os_template.sub(os, i))
        elif i.startswith("@"):
            ret.extend(OS_TEMPLATE[i[1:]])
        else:
            ret.append(i)
    return ret


def replace_arch_template(arch_list: List[str]) -> List[str]:
    ret = []
    for i in arch_list:
        if i.startswith("@"):
            ret.extend(ARCH_TEMPLATE[i[1:]])
        else:
            ret.append(i)
    return ret


def check_and_download(
    url: str,
    dst_file: Path,
    caching=False,
    retries: int = MAX_RETRY,
    display_attempt: int = 1,
    max_attempt: int = 1,
) -> int:
    for net_attempt in range(1, retries + 1):
        try:
            if caching:
                if url in download_cache:
                    logger.info(f"Using cached content: {url}")
                    with dst_file.open("wb") as f:
                        f.write(download_cache[url])
                    return 0
                download_cache[url] = bytes()
            start = time.time()
            with get_session().get(url, stream=True, timeout=(5, 10)) as r:
                r.raise_for_status()
                remote_size = int(r.headers.get("content-length", "0") or "0")
                if "last-modified" in r.headers:
                    remote_ts = parsedate_to_datetime(
                        r.headers["last-modified"]
                    ).timestamp()
                else:
                    remote_ts = None

                written = 0
                last_progress = 0.0
                with dst_file.open("wb") as f:
                    for chunk in r.iter_content(chunk_size=1024**2):
                        if time.time() - start > DOWNLOAD_TIMEOUT:
                            raise TimeoutError("Download timeout")
                        if not chunk:
                            continue  # filter out keep-alive new chunks

                        f.write(chunk)
                        written += len(chunk)
                        if caching:
                            download_cache[url] += chunk
                        if time.time() - last_progress >= 0.2:
                            print_progress(
                                url,
                                written,
                                remote_size,
                                start,
                                status="downloading",
                                attempt=display_attempt,
                                max_attempt=max_attempt,
                            )
                            last_progress = time.time()
                print_progress(
                    url,
                    written,
                    remote_size,
                    start,
                    status="done",
                    attempt=display_attempt,
                    max_attempt=max_attempt,
                    done=True,
                )
                if remote_ts is not None:
                    os.utime(dst_file, (remote_ts, remote_ts))
            return 0
        except BaseException as e:
            logger.error(f"Error occurred (attempt {net_attempt}/{retries}): {e}")
            if dst_file.is_file():
                dst_file.unlink()
            if url in download_cache:
                del download_cache[url]
            if net_attempt < retries:
                sleep_seconds = RETRY_WAIT_SECONDS * net_attempt
                logger.info(f"Retrying {url} in {sleep_seconds:.1f}s")
                print_progress(
                    url,
                    0,
                    0,
                    start_ts=time.time(),
                    status=f"retry in {sleep_seconds:.1f}s",
                    attempt=display_attempt,
                    max_attempt=max_attempt,
                )
                time.sleep(sleep_seconds)
            else:
                print_progress(
                    url,
                    0,
                    0,
                    start_ts=time.time(),
                    status="failed",
                    attempt=display_attempt,
                    max_attempt=max_attempt,
                    done=True,
                )
    return 1


def mkdir_with_dot_tmp(folder: Path) -> Tuple[Path, Path]:
    tmpdir = folder / ".tmp"
    if tmpdir.is_dir():
        shutil.rmtree(str(tmpdir))
    tmpdir.mkdir(parents=True, exist_ok=True)
    return (folder, tmpdir)


def move_files_in(src: Path, dst: Path):
    empty = True
    for file in src.glob("*"):
        empty = False
        logger.info(f"moving {file} to {dst}")
        # shutil.move(str(file), str(dst))
        if file.is_dir():
            (dst / file.name).mkdir(parents=True, exist_ok=True)
            move_files_in(file, dst / file.name)
            file.rmdir()  # rmdir wont fail as all files in it have been moved
        else:
            file.rename(dst / file.name)  # Overwrite files
    if empty:
        logger.info(f"{src} is empty")


def apt_mirror(
    base_url: str,
    dist: str,
    repo: str,
    arch: str,
    dest_base_dir: Path,
    deb_set: Dict[str, int],
    workers: int,
) -> int:
    if not dest_base_dir.is_dir():
        logger.error("Destination directory is empty, cannot continue")
        return 1
    logger.info(f"Started mirroring {base_url} {dist}, {repo}, {arch}!")

    # download Release files
    dist_dir, dist_tmp_dir = mkdir_with_dot_tmp(dest_base_dir / "dists" / dist)
    check_and_download(
        f"{base_url}/dists/{dist}/InRelease", dist_tmp_dir / "InRelease", caching=True
    )
    if (
        check_and_download(
            f"{base_url}/dists/{dist}/Release", dist_tmp_dir / "Release", caching=True
        )
        != 0
    ):
        logger.error("Invalid Repository")
        if not (dist_dir / "Release").is_file():
            logger.warning(
                f"{dist_dir/'Release'} never existed, upstream may not provide packages for {dist}, ignore this error"
            )
            return 0
        return 1
    check_and_download(
        f"{base_url}/dists/{dist}/Release.gpg",
        dist_tmp_dir / "Release.gpg",
        caching=True,
    )

    comp_dir, comp_tmp_dir = mkdir_with_dot_tmp(dist_dir / repo)

    # load Package Index URLs from the Release file
    release_file = dist_tmp_dir / "Release"
    arch_dir = arch if arch in ARCH_NO_PKGIDX else f"binary-{arch}"
    pkgidx_dir, pkgidx_tmp_dir = mkdir_with_dot_tmp(comp_dir / arch_dir)
    with open(release_file, "r") as fd:
        pkgidx_content = None
        cnt_start = False
        for line in fd:
            if cnt_start:
                fields = line.split()
                if (
                    len(fields) != 3 or len(fields[0]) != 64
                ):  # 64 is SHA-256 checksum length
                    break
                checksum, filesize, filename = tuple(fields)
                if (
                    filename.startswith(f"{repo}/{arch_dir}/")
                    or filename.startswith(f"{repo}/Contents-{arch}")
                    or filename.startswith(f"Contents-{arch}")
                ):
                    fn = Path(filename)
                    if len(fn.parts) <= 3:
                        # Contents-amd64.gz
                        # main/Contents-amd64.gz
                        # main/binary-all/Packages
                        pkgidx_file = dist_dir / fn.parent / ".tmp" / fn.name
                    else:
                        # main/dep11/by-hash/MD5Sum/0af5c69679a24671cfd7579095a9cb5e
                        # deep_tmp_dir is in pkgidx_tmp_dir hence no extra garbage collection needed
                        deep_tmp_dir = (
                            dist_dir
                            / Path(fn.parts[0])
                            / Path(fn.parts[1])
                            / ".tmp"
                            / Path("/".join(fn.parts[2:-1]))
                        )
                        deep_tmp_dir.mkdir(parents=True, exist_ok=True)
                        pkgidx_file = deep_tmp_dir / fn.name
                else:
                    logger.warning(f"Ignore the file {filename}")
                    continue
                pkglist_url = f"{base_url}/dists/{dist}/{filename}"
                if check_and_download(pkglist_url, pkgidx_file) != 0:
                    logger.error(f"Failed to download: {pkglist_url}")
                    continue

                with pkgidx_file.open("rb") as t:
                    content = t.read()
                if len(content) != int(filesize):
                    logger.error(f"Invalid size of {pkgidx_file}, expected {filesize}, skipped")
                    pkgidx_file.unlink()
                    continue
                if hashlib.sha256(content).hexdigest() != checksum:
                    logger.error(f"Invalid checksum of {pkgidx_file}, expected {checksum}, skipped")
                    pkgidx_file.unlink()
                    continue
                if pkgidx_content is None and pkgidx_file.stem == "Packages":
                    logger.info(f"getting packages index content from {pkgidx_file.name}")
                    suffix = pkgidx_file.suffix
                    if suffix == ".xz":
                        pkgidx_content = lzma.decompress(content).decode("utf-8")
                    elif suffix == ".bz2":
                        pkgidx_content = bz2.decompress(content).decode("utf-8")
                    elif suffix == ".gz":
                        pkgidx_content = gzip.decompress(content).decode("utf-8")
                    elif suffix == "":
                        pkgidx_content = content.decode("utf-8")
                    else:
                        logger.error("unsupported format")

            # Currently only support SHA-256 checksum, because
            # "Clients may not use the MD5Sum and SHA1 fields for security purposes, and must require a SHA256 or a SHA512 field."
            # from https://wiki.debian.org/DebianRepository/Format#A.22Release.22_files
            if line.startswith("SHA256:"):
                cnt_start = True
    if not cnt_start:
        logger.error("Cannot find SHA-256 checksum")
        return 1

    def collect_tmp_dir():
        try:
            move_files_in(pkgidx_tmp_dir, pkgidx_dir)
            move_files_in(comp_tmp_dir, comp_dir)
            move_files_in(dist_tmp_dir, dist_dir)

            pkgidx_tmp_dir.rmdir()
            comp_tmp_dir.rmdir()
            dist_tmp_dir.rmdir()
            return 0
        except:
            traceback.print_exc()
            return 1

    if arch in ARCH_NO_PKGIDX:
        if collect_tmp_dir() == 1:
            return 1
        logger.info(f"Mirroring {base_url} {dist}, {repo}, {arch} done!")
        return 0

    if pkgidx_content is None:
        logger.error("index is empty, failed")
        if len(list(pkgidx_dir.glob("Packages*"))) == 0:
            logger.warning(
                f"{pkgidx_dir/'Packages'} never existed, upstream may not provide {dist}/{repo}/{arch}, ignore this error"
            )
            return 0
        return 1

    # Download packages
    err = 0
    packages = []
    for pkg in pkgidx_content.split("\n\n"):
        if len(pkg) < 10:  # ignore blanks
            continue
        try:
            pkg_filename = pattern_package_name.search(pkg).group(1)
            pkg_size = int(pattern_package_size.search(pkg).group(1))
            pkg_checksum = pattern_package_sha256.search(pkg).group(1)
        except:
            logger.error("Failed to parse one package description")
            traceback.print_exc()
            err = 1
            continue
        packages.append((pkg_filename, pkg_size, pkg_checksum))

    deb_count = len(packages)
    deb_size = sum(p[1] for p in packages)

    for pkg_filename, pkg_size, _ in packages:
        dest_filename = dest_base_dir / pkg_filename
        if dest_filename.suffix == ".deb":
            deb_set[str(dest_filename.relative_to(dest_base_dir))] = pkg_size

    def download_package(pkg_filename: str, pkg_size: int, pkg_checksum: str) -> int:
        try:
            dest_filename = dest_base_dir / pkg_filename
            dest_dir = dest_filename.parent
            if not dest_dir.is_dir():
                dest_dir.mkdir(parents=True, exist_ok=True)
            if dest_filename.is_file() and dest_filename.stat().st_size == pkg_size:
                logger.info(f"Skipping {pkg_filename}, size {pkg_size}")
                return 0

            pkg_url = f"{base_url}/{pkg_filename}"
            dest_tmp_filename = dest_filename.with_name("._syncing_." + dest_filename.name)
            for retry in range(1, MAX_RETRY + 1):
                if not is_interactive_shell():
                    logger.info(
                        f"downloading {pkg_url} to {dest_filename} "
                        f"(verify attempt {retry}/{MAX_RETRY})"
                    )
                if (
                    check_and_download(
                        pkg_url,
                        dest_tmp_filename,
                        retries=1,
                        display_attempt=retry,
                        max_attempt=MAX_RETRY,
                    )
                    != 0
                ):
                    continue

                sha = hashlib.sha256()
                with dest_tmp_filename.open("rb") as f:
                    for block in iter(lambda: f.read(1024**2), b""):
                        sha.update(block)
                if sha.hexdigest() != pkg_checksum:
                    logger.error(f"Invalid checksum of {dest_filename}, expected {pkg_checksum}")
                    dest_tmp_filename.unlink()
                    continue
                dest_tmp_filename.rename(dest_filename)
                return 0
            logger.error(f"Failed to download {dest_filename}")
            return 1
        except BaseException:
            traceback.print_exc()
            return 1

    if workers <= 1:
        for pkg_filename, pkg_size, pkg_checksum in packages:
            if download_package(pkg_filename, pkg_size, pkg_checksum) != 0:
                err = 1
    else:
        logger.info(f"Downloading packages with {workers} workers")
        with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
            futures = [
                executor.submit(download_package, pkg_filename, pkg_size, pkg_checksum)
                for pkg_filename, pkg_size, pkg_checksum in packages
            ]
            for future in concurrent.futures.as_completed(futures):
                if future.result() != 0:
                    err = 1

    if collect_tmp_dir() == 1:
        return 1
    logger.info(f"Mirroring {base_url} {dist}, {repo}, {arch} done!")
    logger.info(f"{deb_count} packages, {deb_size} bytes in total")
    return err


def apt_delete_old_debs(dest_base_dir: Path, remote_set: Dict[str, int], dry_run: bool):
    on_disk = set(
        [str(i.relative_to(dest_base_dir)) for i in dest_base_dir.glob("**/*.deb")]
    )
    deleting = on_disk - remote_set.keys()
    # print(on_disk)
    # print(remote_set)
    logger.info(f"Deleting {len(deleting)} packages not in the index{' (dry run)' if dry_run else ''}")
    for i in deleting:
        if dry_run:
            logger.info(f"Will delete {i}")
        else:
            logger.info(f"Deleting {i}")
            (dest_base_dir / i).unlink()


def main():

    parser = argparse.ArgumentParser()
    parser.add_argument("base_url", type=str, help="base URL")
    parser.add_argument("os_version", type=str, nargs="?", default="@debian-current",
                        help="e.g. buster,@ubuntu-lts (default: @debian-current)")
    parser.add_argument("component", type=str, help="e.g. multiverse,contrib")
    parser.add_argument("arch", type=str, nargs="?", default="@common",
                        help="e.g. i386,amd64,@all (default: @common -> amd64,i386,arm64,armhf)")
    parser.add_argument("working_dir", type=Path, help="working directory")
    parser.add_argument(
        "--delete", action="store_true", help="delete unreferenced package files"
    )
    parser.add_argument(
        "--delete-dry-run",
        action="store_true",
        help="print package files to be deleted only",
    )
    parser.add_argument(
        "--proxy",
        type=str,
        default=os.getenv("APT_SYNC_PROXY", ""),
        help="proxy URL for HTTP/HTTPS requests (e.g. http://proxy:8080 or socks5h://proxy:1080); "
             "overrides APT_SYNC_PROXY env var",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=DOWNLOAD_WORKERS,
        help=f"number of concurrent package download workers (default: {DOWNLOAD_WORKERS})",
    )
    args = parser.parse_args()

    if args.workers < 1:
        raise ValueError("workers must be >= 1")

    global SESSION_PROXY
    if args.proxy:
        proxy_url = args.proxy
        if proxy_url.startswith("socks"):
            try:
                import socks  # noqa: F401
            except ImportError:
                logger.warning(
                    "SOCKS proxy requested but PySocks is not installed. "
                    "Install it with: pip install requests[socks]"
                )
        SESSION_PROXY = proxy_url
        logger.info(f"Using proxy: {proxy_url}")

    # generate lists of os codenames
    os_list = args.os_version.split(",")
    check_args("os_version", os_list)
    os_list = replace_os_template(os_list)

    # generate a list of components and archs for each os codename
    def generate_list_for_oses(raw: str, name: str) -> List[List[str]]:
        n_os = len(os_list)
        if ":" in raw:
            # specify os codenames for each component
            lists = []
            for l in raw.split(":"):
                list_for_os = l.split(",")
                check_args(name, list_for_os)
                lists.append(list_for_os)
            assert len(lists) == n_os, f"{name} must be specified for each component"
        else:
            # use same os codenames for all components
            l = raw.split(",")
            check_args(name, l)
            lists = [l] * n_os
        return lists

    component_lists = generate_list_for_oses(args.component, "component")
    arch_lists = [replace_arch_template(al) for al in generate_list_for_oses(args.arch, "arch")]

    logger.info(f"Configuration: {os_list=}, {component_lists=}, {arch_lists=}")

    args.working_dir.mkdir(parents=True, exist_ok=True)
    failed = []
    deb_set = {}

    for dist, arch_list, comp_list in zip(os_list, arch_lists, component_lists):
        for comp in comp_list:
            for arch in arch_list:
                if (
                    apt_mirror(
                        args.base_url,
                        dist,
                        comp,
                        arch,
                        args.working_dir,
                        deb_set=deb_set,
                        workers=args.workers,
                    )
                    != 0
                ):
                    failed.append((dist, comp, arch))
    if len(failed) > 0:
        logger.error(f"Failed APT repos of {args.base_url}: {failed}")
        return
    if args.delete or args.delete_dry_run:
        apt_delete_old_debs(args.working_dir, deb_set, args.delete_dry_run)

    if len(REPO_SIZE_FILE) > 0:
        with open(REPO_SIZE_FILE, "a") as fd:
            total_size = sum(deb_set.values())
            fd.write(f"+{total_size}")


if __name__ == "__main__":
    main()