Add restore-vm (#1)

The parser for restore-vm is a bit simpler because it only has positional
arguments and flags at the beginning like a normal program (no --borg-args
shenanigans). It can actually reuse the ArchiveBuilder class for the restore
because the bind mounts are not read-only, but this doesn't work if the
format of the disk changes.

backup_vm/helpers.py

@@ -0,0 +1,35 @@
+import subprocess
+import os
+
+
+def yes(question, default=None):
+    """Ask the user a yes/no question, optionally with a default answer."""
+    while True:
+        print(question, " (Y/n): " if default else " (y/N): ", end="", file=sys.stderr)
+        answer = input().upper().rstrip("\n")
+        if answer in {"Y", "YES", "1"}:
+            return True
+        elif answer in {"N", "NO", "0"}:
+            return False
+        elif default is not None:
+            return default
+
+
+def grouper(iterable, n):
+    """Collect data into fixed-length chunks or blocks, cutting off remaining elements"""
+    args = [iter(iterable)] * n
+    return zip(*args)
+
+
+def list_entries(archive, properties=["type", "size", "health", "bpath"], passphrases=None):
+    """Wrapper around 'borg list' that returns dicts with keys from '--format'."""
+    env = os.environ.copy()
+    if isinstance(passphrases, str):
+        env["BORG_PASSPHRASE"] = passphrases
+    elif passphrases is not None and archive in passphrases:
+        env["BORG_PASSPHRASE"] = passphrases[archive]
+    format_string = "{NUL}".join("{%s}" % p for p in properties) + "{NUL}"
+    p = subprocess.run(["borg", "list", "--format", format_string, str(archive)],
+                       env=env, stdout=subprocess.PIPE, check=True)
+    for entry in grouper(p.stdout.split(b"\x00"), len(properties)):
+        yield {key: entry[i] for i, key in enumerate(properties)}

backup_vm/parse.py

@@ -402,3 +402,56 @@ def help(self, short=False):
               -p, --progress   force progress display even if stdout isn't a tty
               --borg-args ...  extra arguments passed straight to borg
             """).strip("\n"))
+
+
+class RestoreArgumentParser(ArgumentParser):
+
+    """Argument parser for restore-vm.
+
+    Parses only positional arguments and basic flags (--help, --progress).
+    """
+
+    def __init__(self, default_name="restore-vm", args=sys.argv):
+        self.domain = None
+        super().__init__(default_name, args)
+
+    def parse_args(self, args):
+        positional_args = []
+        for arg in args:
+            if arg in {"-h", "--help"}:
+                self.help()
+                sys.exit()
+            elif arg in {"-p", "--progress"}:
+                self.progress = True
+            elif arg in {"-v", "--version"}:
+                self.version()
+                sys.exit()
+            elif arg.startswith("-"):
+                self.error("unrecognized argument: '{}'".format(arg))
+            else:
+                positional_args.append(arg)
+        try:
+            self.domain, *disks, archive = positional_args
+            self.disks = set(disks)
+            self.archive = Location(archive)
+        except ValueError:
+            self.error("the following arguments are required: domain, archive")
+
+    def help(self, short=False):
+        print(dedent("""
+            usage: {} [-hpv] domain [disk [disk ...]] archive
+        """.format(self.prog).lstrip("\n")))
+        if not short:
+            print(dedent("""
+            Restore a libvirt-based VM from a borg backup.
+
+            positional arguments:
+              domain           libvirt domain to restore
+              disk             a domain block device to restore (default: all disks)
+              archive          a borg archive path (same format as borg create)
+
+            optional arguments:
+              -h, --help       show this help message and exit
+              -v, --version    show version of the backup-vm package
+              -p, --progress   force progress display even if stdout isn't a tty
+            """).strip("\n"))

backup_vm/restore.py

@@ -0,0 +1,72 @@
+import os.path
+import sys
+import libvirt
+from . import lock
+from . import parse
+from . import multi
+from . import builder
+from . import helpers
+from . import snapshot
+
+
+def main():
+    args = parse.RestoreArgumentParser()
+    conn = libvirt.open()
+    if conn is None:
+        print("Failed to open connection to libvirt", file=sys.stderr)
+        sys.exit(1)
+    try:
+        dom = conn.lookupByName(args.domain)
+    except libvirt.libvirtError:
+        print("Domain '{}' not found".format(args.domain))
+        sys.exit(1)
+
+    dom_disks = set(parse.Disk.from_domain(dom))
+    if len(dom_disks) == 0:
+        print("Domain has no disks(!)", file=sys.stderr)
+        sys.exit(1)
+
+    disks_to_restore = args.disks and {x for x in dom_disks if x.target in args.disks} or dom_disks
+    targets_to_restore = {x.target for x in disks_to_restore}
+    if (len(args.disks) > 0 and args.disks != targets_to_restore
+        and not helpers.yes("Some disks to be restored don't exist on the domain: " +
+                    " ".join(sorted(args.disks - targets_to_restore)) +
+                    "\nDo you want to recreate them on the target domain?", False)):
+        sys.exit(1)
+
+    archive_disks = {}
+    passphrases = multi.get_passphrases([args.archive]) if sys.stdout.isatty() else None
+    for entry in helpers.list_entries(args.archive, passphrases=passphrases):
+        if entry["type"] == b"-":
+            name = entry["bpath"].decode("utf-8")
+            if name.split(".")[0] in targets_to_restore:
+                if (entry["health"] != b"healthy"
+                    and not helpers.yes("The backup copy of disk {} is marked as broken.".format(name.split(".")[0]) +
+                                "\nDo you still want to replace it with a possibly corrupted version?", False)):
+                    sys.exit(1)
+                else:
+                    archive_disks[name] = int(entry["size"])
+    if len(disks_to_restore) != len(archive_disks):
+        archive_tgts = {d.split(".")[0] for d in archive_disks.keys()}
+        print("Some disks to be restored don't exist in the archive:",
+              *sorted(targets_to_restore - archive_tgts), file=sys.stderr)
+        sys.exit(1)
+    # TODO: warn if a disk's logical size differs from backup to target
+    # this is complicated by non-raw images like qcow2 (which are expandable)
+
+    with lock.DiskLock(disks_to_restore), builder.ArchiveBuilder(disks_to_restore) as archive_dir:
+        args.archive.extra_args.extend(archive_disks.keys())
+        borg_failed = multi.assimilate(
+            archives=[args.archive],
+            total_size=args.progress and archive_dir.total_size,
+            dir_to_archive=None,
+            passphrases=passphrases,
+            verb="extract"
+        )
+
+    # bug in libvirt python wrapper(?): sometimes it tries to delete
+    # the connection object before the domain, which references it
+    del dom
+    del conn
+
+    sys.exit(borg_failed)

setup.py

@@ -37,6 +37,7 @@ def readme():
       entry_points={
           "console_scripts": [
               "backup-vm=backup_vm.backup:main",
+              "restore-vm=backup_vm.restore:main",
               "borg-multi=backup_vm.multi:main",
           ],
       },