Refactor into package

MANIFEST.in

@@ -0,0 +1 @@
+include README.rst

backup_vm/backup.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+
+import os.path
+import sys
+import libvirt
+from . import parse
+from . import multi
+from . import builder
+from . import snapshot
+
+
+def main():
+    args = parse.ArgumentParser(sys.argv)
+    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)
+
+    if args.memory and not dom.isActive():
+        print("Domain is shut off, cannot save memory contents", file=sys.stderr)
+        args.memory = False
+
+    all_disks = set(parse.Disk.get_disks(dom))
+    if len(all_disks) == 0:
+        print("Domain has no disks(!)", file=sys.stderr)
+        sys.exit(1)
+
+    disks_to_backup = args.disks and {x for x in all_disks if x.target in args.disks} or all_disks
+    if len(disks_to_backup) != len(args.disks or all_disks):
+        print("Some disks to be backed up don't exist on the domain:",
+              *sorted(args.disks - {x.target for x in disks}), file=sys.stderr)
+        sys.exit(1)
+
+    for disk in all_disks:
+        disk.failed = False
+        if disk not in disks_to_backup:
+            disk.snapshot_path = None
+            continue
+        filename = args.domain + "-" + disk.target + "-tempsnap.qcow2"
+        if disk.type == "dev":
+            # we probably can't write the temporary snapshot to the same directory
+            # as the original disk, so use the default libvirt images directory
+            disk.snapshot_path = os.path.join("/var/lib/libvirt/images", filename)
+        else:
+            disk.snapshot_path = os.path.join(os.path.dirname(disk.path), filename)
+
+    memory = os.path.join(tmpdir, "memory.bin") if args.memory else None
+    with snapshot.Snapshot(dom, all_disks, memory, args.progress), \
+            builder.ArchiveBuilder(disks_to_backup) as archive_dir:
+        if args.progress:
+            borg_failed = multi.assimilate(args.archives, archive_dir.total_size)
+        else:
+            borg_failed = multi.assimilate(args.archives)
+        pass
+
+    # 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 or any(disk.failed for disk in disks_to_backup))

backup_vm/builder.py

@@ -0,0 +1,47 @@
+import subprocess
+import tempfile
+import os.path
+
+
+class ArchiveBuilder(tempfile.TemporaryDirectory):
+
+    """Creates the folder to be turned into a VM backup.
+
+    Creates a temporary folder populated with symlinks to each disk to backup.
+    Essentially lays out the contents of the archive to be created.
+
+    Attributes:
+        name: The path of the temporary directory.
+        total_size: The total size of every disk linked to in the directory.
+    """
+
+    def __init__(self, disks, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.total_size = 0
+        self.disks = disks
+        self.old_cwd = os.getcwd()
+        os.chdir(self.name)
+
+    def __enter__(self):
+        for disk in self.disks:
+            realpath = os.path.realpath(disk.path)
+            with open(realpath) as f:
+                # add size of disk to total
+                f.seek(0, os.SEEK_END)
+                self.total_size += f.tell()
+            linkpath = disk.target + "." + disk.format
+            with open(linkpath, "w") as f:
+                # simulate 'touch'
+                pass
+            # following symlinks for --read-special is still broken :(
+            # when issue gets fixed should switch to symlinks:
+            # https://github.com/borgbackup/borg/issues/1215
+            subprocess.run(["mount", "--bind", realpath, linkpath], check=True)
+        return self
+
+    def cleanup(self):
+        for disk in self.disks:
+            linkpath = disk.target + "." + disk.format
+            subprocess.run(["umount", linkpath], check=True)
+        os.chdir(self.old_cwd)
+        return super().cleanup()

backup_vm/multi.py

@@ -0,0 +1,236 @@
+from distutils.version import LooseVersion
+from base64 import b64encode
+from getpass import getpass
+from pty import openpty
+from copy import copy
+import subprocess
+import selectors
+import termios
+import fcntl
+import json
+import sys
+import pty
+import os
+
+
+def get_passphrases(archives):
+    """Prompts the user for their archive passphrases.
+
+    Checks for archives that won't open without a (non-blank, non-random)
+    BORG_PASSPHRASE and prompts the user for their passphrases.
+
+    Args:
+        archives: A list of Location objects to check the repositories of.
+
+    Returns:
+        A dictionary mapping archives to their (purported) passphrases. The
+        entered passphrases are not checked to actually open the archives.
+    """
+    passphrases = {}
+    env = os.environ.copy()
+    for archive in archives:
+        repo = copy(archive)
+        repo.archive = None
+        # check if we need a password as recommended by the docs:
+        # https://borgbackup.readthedocs.io/en/stable/internals/frontends.html#passphrase-prompts
+        if len({"BORG_PASSPHRASE", "BORG_PASSCOMMAND", "BORG_NEWPASSPHRASE"} - set(env)) == 3:
+            # generate random password that would be incorrect were it needed
+            env["BORG_PASSPHRASE"] = b64encode(os.urandom(16)).decode("utf-8")
+        with subprocess.Popen(["borg", "list", str(repo)], stdin=subprocess.PIPE,
+                              stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, env=env) as proc:
+            # manually close stdin instead of /dev/null so borg knows it won't get input
+            proc.stdin.close()
+            proc.stdin = None
+            err = proc.communicate(input)[1].decode("utf-8").rstrip("\n").split("\n")[-1]
+            if proc.poll() != 0:
+                # exact error message changes between borg versions
+                if err.startswith("passphrase supplied") and err.endswith("is incorrect."):
+                    passphrases[archive] = getpass("Enter passphrase for key {!s}: ".format(repo))
+    return passphrases
+
+
+def log(name, msg, *args, file=sys.stderr, end="\n", **kwargs):
+    """Logs a string to a file, prepending a "tag" to each line.
+
+    Logs a string to a file (by default stderr), with a "tag" added to the
+    beginning of each line, in the format of this example::
+
+        [repo::archive] Hello world!
+
+    Args:
+        name: The text to be put in the "tag" part of each line.
+        msg: The string to be tagged & logged.
+        end: The ending of the last line printed.
+
+    Any other arguments passed will be passed onto print().
+    """
+    for l in msg[:-1]:
+        print("[{}] {}".format(name, l), file=file, **kwargs)
+    print("[{}] {}".format(name, msg[-1]), file=file, end=end, **kwargs)
+
+
+def process_line(p, line, total_size=None, prompt_answers={}):
+    """Process a line coming from a borg process.
+
+    Processes JSON emitted by a borg process with --log-json turned on. The
+    lines are cached, so 1 line does not have to equal 1 JSON message.
+
+    Args:
+        p: The process the line came from (with some extra properties added to
+            the Popen object).
+        line: The line read from the process's stdout or stderr. If it contains
+            progress information, update the stored progress value. If it is a
+            prompt for the user, ask for and return the answer (& cache it for
+            later.) If it is a log message or some other non-JSON, print it out.
+        total_size: The total size of all files being backed up. This can be set
+            to None to disable progress calculation.
+        prompt_answers: A dictionary of previous answers from users' prompts.
+            Prompts with msgids in the dictionary will be automatically answered
+            with the value given (ostensibly from an earlier prompt).
+    """
+    if len(p.json_buf) > 0 or line.startswith("{"):
+        p.json_buf.append(line)
+    if len(p.json_buf) > 0 and line.endswith("}"):
+        try:
+            msg = json.loads("\n".join(p.json_buf))
+            p.json_buf = []
+            if msg["type"] == "archive_progress" and total_size is not None:
+                p.progress = msg["original_size"] / total_size
+            elif msg["type"] == "log_message":
+                log(p.archive.orig, msg["message"].split("\n"))
+            elif msg["type"].startswith("question"):
+                if "msgid" in msg:
+                    prompt_id = msg["msgid"]
+                elif "message" in msg:
+                    prompt_id = msg["message"]
+                else:
+                    raise ValueError("No msgid or message for prompt")
+                if msg.get("is_prompt", False) or msg["type"].startswith("question_prompt"):
+                    if prompt_id not in prompt_answers:
+                        log(p.archive.orig, msg["message"].split("\n"), end="")
+                        try:
+                            prompt_answers[prompt_id] = input()
+                            print(prompt_answers[prompt_id], file=p.stdin, flush=True)
+                        except EOFError:
+                            p.stdin.close()
+                elif not msg["type"].startswith("question_accepted"):
+                    log(p.archive.orig, msg["message"].split("\n"))
+        except json.decoder.JSONDecodeError:
+            log(p.archive.orig, p.json_buf)
+            p.json_buf = []
+    elif line.startswith("Enter passphrase for key "):
+        log(p.archive.orig, [line], end="")
+        passphrase = getpass("")
+        print(passphrase, file=p.stdin, flush=True)
+        print("", file=sys.stderr)
+    elif line != "":
+        # line is not json?
+        log(p.archive.orig, [line])
+    # TODO: process password here for efficiency & simplicity
+
+
+def get_borg_version():
+    """
+    Get the version of the system borg.
+
+    Returns:
+        The version of the system borg as a distutils.version.LooseVersion (for
+        easy comparison with other versions).
+    """
+    version_bytes = subprocess.run(["borg", "--version"], stdout=subprocess.PIPE, check=True).stdout
+    return LooseVersion(version_bytes.decode("utf-8").split(" ")[1])
+
+
+def assimilate(archives, total_size=None, dir_to_archive=".", passphrases=None):
+    """
+    Run and manage multiple `borg create` commands.
+
+    Args:
+        archives: A list containing Location objects for the archives to create.
+        total_size: The total size of all files being backed up. As borg
+            normally only makes one pass over the data, it can't calculate
+            percentages on its own. Setting this to None disables progress
+            calculation.
+        dir_to_archive: The directory to archive. Defaults to the current
+            directory.
+
+    Returns:
+        A boolean indicating if any borg processes failed (True = failed).
+    """
+
+    if passphrases is None and sys.stdout.isatty():
+        passphrases = get_passphrases(archives)
+
+    if get_borg_version() < LooseVersion("1.1.0"):
+        # borg <1.1 doesn't support --log-json for the progress display
+        print("You are using an old version of borg, progress indication is disabled", file=sys.stderr)
+        recent_borg = False
+        progress = False
+    else:
+        recent_borg = True
+        progress = total_size is not None
+
+    borg_processes = []
+    borg_failed = False
+    try:
+        with selectors.DefaultSelector() as sel:
+            for idx, archive in enumerate(archives):
+                if progress:
+                    archive.extra_args.append("--progress")
+                if recent_borg:
+                    archive.extra_args.append("--log-json")
+                env = os.environ.copy()
+                passphrase = passphrases.get(archive, os.environ.get("BORG_PASSPHRASE"))
+                if passphrase is not None:
+                    env["BORG_PASSPHRASE"] = passphrase
+                master, slave = openpty()
+                settings = termios.tcgetattr(master)
+                settings[3] &= ~termios.ECHO
+                termios.tcsetattr(master, termios.TCSADRAIN, settings)
+                proc = subprocess.Popen(["borg", "create", str(archive), dir_to_archive, "--read-special", *archive.extra_args],
+                                        stdout=slave, stderr=slave, stdin=slave, close_fds=True, env=env, start_new_session=True)
+                fl = fcntl.fcntl(master, fcntl.F_GETFL)
+                fcntl.fcntl(master, fcntl.F_SETFL, fl | os.O_NONBLOCK)
+                proc.stdin = os.fdopen(master, "w")
+                proc.stdout = os.fdopen(master, "r")
+                proc.archive = archive
+                proc.json_buf = []
+                proc.progress = 0
+                borg_processes.append(proc)
+                sel.register(proc.stdout, selectors.EVENT_READ, data=proc)
+
+            if progress:
+                print("backup progress: 0%".ljust(25), end="\u001b[25D", flush=True)
+            else:
+                # give the user some feedback so the program doesn't look frozen
+                print("starting backup", flush=True)
+            while len(sel.get_map()) > 0:
+                for key, mask in sel.select(1):
+                    for line in iter(key.fileobj.readline, ""):
+                        process_line(key.data, line.rstrip("\n"), total_size)
+                for key in [*sel.get_map().values()]:
+                    if key.data.poll() is not None:
+                        key.data.wait()
+                        key.data.progress = 1
+                        if key.data.returncode != 0:
+                            borg_failed = True
+                        sel.unregister(key.fileobj)
+                if progress:
+                    total_progress = sum(p.progress for p in borg_processes)
+                    print("backup progress: {}%".format(
+                        int(total_progress / len(borg_processes) * 100)).ljust(25), end="\u001b[25D")
+            if progress:
+                print()
+    finally:
+        for p in borg_processes:
+            if p.poll() is not None:
+                p.kill()
+                try:
+                    p.communicate()
+                except (ValueError, OSError):
+                    p.wait()
+    return borg_failed
+
+
+def main():
+    pass