#!/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.
"""

import argparse
import datetime
import json
import os
import shutil
import subprocess
import sys
import time
import urllib.error
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 parse_args():
    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=datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
                        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")
    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")
    parser.add_argument("--profiles", default="", help="compose profiles to enable, comma-separated")
    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")
    parser.add_argument("--services-exclude", default="", help="services to exclude, comma-separated")
    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):
        self.handle = handle
        self.console = console

    def write(self, text):
        self.console.write(text)
        self.handle.write(text)
        self.handle.flush()

    def flush(self):
        self.console.flush()


def fail(message):
    print("error - " + message, file=sys.stderr)
    sys.exit(1)


def run(command, cwd=None):
    """Run a command, fail on errors and return its output."""
    result = subprocess.run(command, cwd=cwd, check=False, text=True, capture_output=True)
    if result.returncode != 0:
        fail("'%s' failed:\n%s" % (" ".join(command), result.stderr.strip()))
    return result.stdout


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

    def __init__(self, args):
        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("skip service (excluded) '%s'" % 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 = "volume_" + source
            name = self.config.get("volumes", {}).get(source, {}).get("name") or "%s_%s" % (self.config["name"], source)
            inspect = subprocess.run(["docker", "volume", "inspect", "--format", "{{ .Mountpoint }}", name],
                                     check=False, text=True, capture_output=True)
            if inspect.returncode != 0:
                print("skip volume (not created) '%s'" % name)
                return None
            path = inspect.stdout.strip()
        else:
            return None
        mount = {"type": kind, "source": source, "label": label.replace("/", "_"), "path": path}
        if not os.path.isdir(path):
            print("skip volume (not a directory) '%s'" % path)
        elif os.path.ismount(path):
            print("skip volume (mount point) '%s'" % path)
        else:
            print("use volume '%s'" % path)
            return mount
        return None

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

    def down(self):
        print("shutdown services")
        run(self.compose + ["down"], cwd=self.compose_dir)

    def up(self, services):
        if not services:
            print("no services were running, skip start")
            return
        print("start services '%s'" % " ".join(services))
        run(self.compose + ["up", "--detach"] + services, cwd=self.compose_dir)


def create(args, project):
    manifest_path = os.path.join(args.backup_dir, args.backup_name + ".manifest.json")
    suffix = ".tar.gz" if args.compress else ".tar"
    if os.path.exists(manifest_path):
        fail("backup '%s' exists already" % args.backup_name)
    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"] = "%s_backup_%s%s" % (args.backup_name, 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("create backup from volume '%s' to archive '%s'" % (mount["path"], archive))
            run(["tar", "--create", "--auto-compress", "--file", archive + ".tmp", "--directory", mount["path"], "."])
            os.replace(archive + ".tmp", archive)
        with open(manifest_path + ".tmp", "w", encoding="utf-8") as handle:
            json.dump(manifest, handle, indent=2)
        os.replace(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("remove backup file (retention) '%s'" % path)
            os.remove(path)


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


def restore(args, project):
    manifest_path = os.path.join(args.backup_dir, args.backup_name + ".manifest.json")
    if not os.path.isfile(manifest_path):
        fail("manifest '%s' not found" % manifest_path)
    with open(manifest_path, encoding="utf-8") as handle:
        manifest = json.load(handle)
    if manifest.get("version") != MANIFEST_VERSION:
        fail("manifest '%s' has an unsupported version" % manifest_path)
    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("volume '%s' of the backup is not part of the project" % mount["source"])
        if not os.path.isfile(archive):
            fail("archive '%s' not found" % archive)
        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("restore backup from archive '%s' to volume '%s'" % (archive, 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():
    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 = "%s.create.log" % args.backup_name
    else:
        log = "%s.restore_%s.log" % (args.backup_name, time.strftime("%Y-%m-%d_%H-%M-%S"))
    handle = open(os.path.join(args.backup_dir, log), "a", encoding="utf-8")  # pylint: disable=consider-using-with
    sys.stdout, sys.stderr = Log(handle, sys.stdout), Log(handle, sys.stderr)
    print("start time: %s" % datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"))
    print("use backup name '%s'" % args.backup_name)
    print("use backup directory '%s'" % args.backup_dir)
    print("use compose directory '%s'" % args.compose_dir)
    project = Project(args)
    if args.create:
        create(args, project)
    else:
        restore(args, project)
    print("end time: %s" % datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"))


if __name__ == "__main__":
    main()
