#!/usr/bin/env python3
"""Back up and restore the volumes of a Docker Compose project.

The volumes are taken from `docker compose config`, so Compose resolves the
compose files, variables and profiles. Directories bound into containers and
named volumes are archived with tar while the project is down, files, sockets
and mount points are skipped. Every backup consists of one archive per volume
and a manifest, the restore checks both before it stops anything. The output
of every run is also written to a log file next to the manifest.

Run as `dc-backup --create` or `dc-backup --restore` with the options below;
the role docker_compose installs the script and runs `--create` from cron.
"""

import argparse
import datetime
import json
import os
import shutil
import subprocess
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

MANIFEST_VERSION = 1
# Cron starts with a minimal PATH, so the usual directories are appended.
DEFAULT_PATH = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"


def timestamp():
    """Return the current date and time in the form the backup names use."""
    return datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")


def parse_args():
    """Parse the command line; the module docstring is the help text."""
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    action = parser.add_mutually_exclusive_group(required=True)
    action.add_argument(
        "--create", action="store_true", help="create a backup of the volumes"
    )
    action.add_argument(
        "--restore", action="store_true", help="restore a backup of the volumes"
    )
    parser.add_argument(
        "--backup-dir",
        default="./backups",
        help="backup directory, defaults to './backups'",
    )
    parser.add_argument(
        "--backup-name",
        default=timestamp(),
        help="backup name, defaults to the current date and time",
    )
    parser.add_argument(
        "--compose-dir", default=".", help="compose directory, defaults to '.'"
    )
    parser.add_argument(
        "--compose-files",
        default="docker-compose.yml",
        help="compose files relative to the compose directory, comma-separated, "
        "defaults to 'docker-compose.yml'",
    )
    parser.add_argument(
        "--compress", action="store_true", help="compress the archives with gzip"
    )
    parser.add_argument(
        "--heartbeat-url",
        default="",
        help="URL to request after a successful backup, defaults to none",
    )
    parser.add_argument(
        "--profiles",
        default="",
        help="compose profiles to enable, comma-separated, defaults to none",
    )
    parser.add_argument(
        "--retention-days",
        type=int,
        default=0,
        help="remove files of the backup directory which are this many days old "
        "after a successful backup, defaults to 0, which keeps everything",
    )
    parser.add_argument(
        "--services-exclude",
        default="",
        help="services to exclude, comma-separated, defaults to none",
    )
    args = parser.parse_args()
    args.backup_dir = os.path.abspath(args.backup_dir)
    args.compose_dir = os.path.abspath(args.compose_dir)
    args.compose_files = [
        os.path.join(args.compose_dir, f) for f in args.compose_files.split(",") if f
    ]
    args.profiles = [p for p in args.profiles.split(",") if p]
    args.services_exclude = [s for s in args.services_exclude.split(",") if s]
    return args


class Log:
    """Writes the output of a console stream to the log file of the run as well."""

    def __init__(self, handle, console):
        """Wrap the console stream and the open log file."""
        self.handle = handle
        self.console = console

    def write(self, text):
        """Write the text to the console and to the log file."""
        self.console.write(text)
        self.handle.write(text)
        self.handle.flush()

    def flush(self):
        """Flush the console; the log file is flushed on every write."""
        self.console.flush()


def fail(message):
    """Print the message as an error on stderr and exit with status 1."""
    print(f"error - {message}", file=sys.stderr)
    sys.exit(1)


def run(command, cwd=None, check=True):
    """Run a command and return its output.

    A failure ends the script, unless check is False: then None is returned,
    for a command whose failure is an expected answer.
    """
    result = subprocess.run(
        command, cwd=cwd, check=False, text=True, capture_output=True
    )
    if result.returncode != 0:
        if not check:
            return None
        joined = " ".join(command)
        fail(f"'{joined}' failed:\n{result.stderr.strip()}")
    return result.stdout


class Project:
    """The Compose project with its volumes."""

    def __init__(self, args):
        """Resolve the Compose configuration and collect the mounts."""
        self.compose_dir = args.compose_dir
        self.compose = ["docker", "compose"] + [
            a for f in args.compose_files for a in ("--file", f)
        ]
        self.compose += [a for p in args.profiles for a in ("--profile", p)]
        self.config = json.loads(
            run(self.compose + ["config", "--format", "json"], cwd=self.compose_dir)
        )
        self.mounts = self.find_mounts(args.services_exclude)

    def find_mounts(self, services_exclude):
        """Return the mounts to archive, one entry per directory."""
        mounts = {}
        for service, definition in sorted(self.config.get("services", {}).items()):
            if service in services_exclude:
                print(f"skip service (excluded) '{service}'")
                continue
            for volume in definition.get("volumes", []):
                mount = self.resolve_mount(volume)
                if mount is not None and mount["path"] not in mounts:
                    mounts[mount["path"]] = mount
        return list(mounts.values())

    def resolve_mount(self, volume):
        """Return the label and path of a volume, or None when it is skipped."""
        kind, source = volume.get("type"), volume.get("source")
        if kind == "bind":
            label = os.path.relpath(source, self.compose_dir)
            if label.startswith(".."):
                label = source.lstrip("/")
            path = source
        elif kind == "volume" and source:
            label = f"volume_{source}"
            name = (
                self.config.get("volumes", {}).get(source, {}).get("name")
                or f"{self.config['name']}_{source}"
            )
            mountpoint = run(
                ["docker", "volume", "inspect", "--format", "{{ .Mountpoint }}", name],
                check=False,
            )
            if mountpoint is None:
                print(f"skip volume (not created) '{name}'")
                return None
            path = mountpoint.strip()
        else:
            return None
        mount = {
            "type": kind,
            "source": source,
            "label": label.replace("/", "_"),
            "path": path,
        }
        if not os.path.isdir(path):
            print(f"skip volume (not a directory) '{path}'")
        elif os.path.ismount(path):
            print(f"skip volume (mount point) '{path}'")
        else:
            print(f"use volume '{path}'")
            return mount
        return None

    def running_services(self):
        """Return the names of the services that are running."""
        return run(
            self.compose + ["ps", "--status", "running", "--services"],
            cwd=self.compose_dir,
        ).split()

    def down(self):
        """Stop and remove the containers of the project."""
        print("shutdown services")
        run(self.compose + ["down"], cwd=self.compose_dir)

    def up(self, services):
        """Start the given services again; nothing when none was running."""
        if not services:
            print("no services were running, skip start")
            return
        joined = " ".join(services)
        print(f"start services '{joined}'")
        run(self.compose + ["up", "--detach"] + services, cwd=self.compose_dir)


def create(args, project):
    """Archive every mount of the project while it is down and write the manifest."""
    manifest_path = os.path.join(args.backup_dir, f"{args.backup_name}.manifest.json")
    suffix = ".tar.gz" if args.compress else ".tar"
    if os.path.exists(manifest_path):
        fail(f"backup '{args.backup_name}' exists already")
    if not project.mounts:
        fail("no volumes to back up")
    manifest = {"version": MANIFEST_VERSION, "name": args.backup_name, "mounts": []}
    for mount in project.mounts:
        mount["archive"] = f"{args.backup_name}_backup_{mount['label']}{suffix}"
        manifest["mounts"].append(mount)
    os.makedirs(args.backup_dir, exist_ok=True)
    running = project.running_services()
    project.down()
    try:
        for mount in project.mounts:
            archive = os.path.join(args.backup_dir, mount["archive"])
            print(f"create backup from volume '{mount['path']}' to archive '{archive}'")
            run(
                [
                    "tar",
                    "--create",
                    "--auto-compress",
                    "--file",
                    f"{archive}.tmp",
                    "--directory",
                    mount["path"],
                    ".",
                ]
            )
            os.replace(f"{archive}.tmp", archive)
        with open(f"{manifest_path}.tmp", "w", encoding="utf-8") as handle:
            json.dump(manifest, handle, indent=2)
        os.replace(f"{manifest_path}.tmp", manifest_path)
    finally:
        for name in os.listdir(args.backup_dir):
            if name.startswith(args.backup_name) and name.endswith(".tmp"):
                os.remove(os.path.join(args.backup_dir, name))
        project.up(running)
    apply_retention(args)
    send_heartbeat(args)


def apply_retention(args):
    """Remove the files of earlier backups, which are named after their date."""
    if args.retention_days <= 0:
        return
    for name in sorted(os.listdir(args.backup_dir)):
        path = os.path.join(args.backup_dir, name)
        if (
            name.startswith("20")
            and os.path.isfile(path)
            and time.time() - os.path.getmtime(path) >= args.retention_days * 86400
        ):
            print(f"remove backup file (retention) '{path}'")
            os.remove(path)


def masked_url(url):
    """Return the scheme and host of the URL; path and credentials stay out."""
    parts = urllib.parse.urlsplit(url)
    host = parts.hostname or ""
    if parts.port is not None:
        host = f"{host}:{parts.port}"
    return f"{parts.scheme}://{host}/..."


def send_heartbeat(args):
    """Request the heartbeat URL like `wget --spider`, a failure is only reported."""
    if not args.heartbeat_url:
        return
    shown = masked_url(args.heartbeat_url)
    for method in ("HEAD", "GET"):
        try:
            urllib.request.urlopen(
                urllib.request.Request(args.heartbeat_url, method=method), timeout=30
            )
            print(f"sent heartbeat to '{shown}'")
            return
        except urllib.error.HTTPError as error:
            message = f"{error.code} {error.reason}"
        except OSError as error:
            message = str(error)
    print(f"warning - heartbeat to '{shown}' failed: {message}", file=sys.stderr)


def restore(args, project):
    """Check the backup, then replace the content of every mount while down."""
    manifest_path = os.path.join(args.backup_dir, f"{args.backup_name}.manifest.json")
    if not os.path.isfile(manifest_path):
        fail(f"manifest '{manifest_path}' not found")
    with open(manifest_path, encoding="utf-8") as handle:
        manifest = json.load(handle)
    if manifest.get("version") != MANIFEST_VERSION:
        fail(f"manifest '{manifest_path}' has an unsupported version")
    current = {(m["type"], m["source"]): m["path"] for m in project.mounts}
    jobs = []
    for mount in manifest["mounts"]:
        archive = os.path.join(args.backup_dir, mount["archive"])
        if (mount["type"], mount["source"]) not in current:
            fail(f"volume '{mount['source']}' of the backup is not part of the project")
        if not os.path.isfile(archive):
            fail(f"archive '{archive}' not found")
        run(["tar", "--list", "--file", archive])
        jobs.append((archive, current[(mount["type"], mount["source"])]))
    running = project.running_services()
    project.down()
    try:
        for archive, path in jobs:
            print(f"restore backup from archive '{archive}' to volume '{path}'")
            for entry in os.scandir(path):
                if entry.is_dir(follow_symlinks=False):
                    shutil.rmtree(entry.path)
                else:
                    os.remove(entry.path)
            run(["tar", "--extract", "--file", archive, "--directory", path])
    finally:
        project.up(running)


def main():
    """Run the backup or the restore with every line of output logged."""
    args = parse_args()
    os.environ["PATH"] = os.environ.get("PATH", "") + os.pathsep + DEFAULT_PATH
    if shutil.which("docker") is None:
        fail("docker not found in PATH")
    os.makedirs(args.backup_dir, exist_ok=True)
    if args.create:
        log = f"{args.backup_name}.create.log"
    else:
        log = f"{args.backup_name}.restore_{timestamp()}.log"
    # The log stays open for the whole run and closes with the process.
    handle = open(  # noqa: SIM115
        os.path.join(args.backup_dir, log), "a", encoding="utf-8"
    )
    sys.stdout, sys.stderr = Log(handle, sys.stdout), Log(handle, sys.stderr)
    print(f"start time: {timestamp()}")
    print(f"use backup name '{args.backup_name}'")
    print(f"use backup directory '{args.backup_dir}'")
    print(f"use compose directory '{args.compose_dir}'")
    project = Project(args)
    if args.create:
        create(args, project)
    else:
        restore(args, project)
    print(f"end time: {timestamp()}")


if __name__ == "__main__":
    main()
