backupbot.py

@@ -9,7 +9,6 @@ import logging
 import docker
 import restic
 import tarfile
-import io
 from pythonjsonlogger import jsonlogger
 from datetime import datetime, timezone
 from restic.errors import ResticFailedError
@@ -19,6 +18,8 @@ from shutil import copyfile, rmtree
 VOLUME_PATH = "/var/lib/docker/volumes/"
 SECRET_PATH = "/secrets/"
 SERVICE = "ALL"
+BACKUP_FILE = "/tmp/backup.tar.gz"
+DOWNLOAD_TMP_PATH = "/tmp/backup-download"
 
 logger = logging.getLogger("backupbot")
 logging.addLevelName(55, "SUMMARY")
@@ -573,53 +574,85 @@ def download(snapshot, path, volumes, secrets):
         snapshot = latest_snapshot[0]["short_id"]
     if not any([path, volumes, secrets]):
         volumes = secrets = True
-    if path:
-        path = path.removesuffix("/")
-        binary_output = dump(snapshot, path)
-        files = list_files(snapshot, path)
-        filetype = [f.get("type") for f in files if f.get("path") == path][0]
-        filename = Path(path).name
-        if filetype == "dir":
-            filename = filename + ".tar"
-        tarinfo = tarfile.TarInfo(name=filename)
-        tarinfo.size = len(binary_output)
-        file_dumps.append((binary_output, tarinfo))
-    if volumes:
-        if SERVICE == "ALL":
-            logger.error("Please specify '--host' when using '--volumes'")
-            exit(1)
-        files = list_files(snapshot, VOLUME_PATH)
-        for f in files[1:]:
-            path = f["path"]
-            if Path(path).name.startswith(SERVICE) and f["type"] == "dir":
-                binary_output = dump(snapshot, path)
-                filename = f"{Path(path).name}.tar"
-                tarinfo = tarfile.TarInfo(name=filename)
-                tarinfo.size = len(binary_output)
-                file_dumps.append((binary_output, tarinfo))
-    if secrets:
-        if SERVICE == "ALL":
-            logger.error("Please specify '--host' when using '--secrets'")
-            exit(1)
-        filename = f"{SERVICE}.json"
-        files = list_files(snapshot, SECRET_PATH)
-        secrets = {}
-        for f in files[1:]:
-            path = f["path"]
-            if Path(path).name.startswith(SERVICE) and f["type"] == "file":
-                secret = dump(snapshot, path).decode()
-                secret_name = path.removeprefix(f"{SECRET_PATH}{SERVICE}_")
-                secrets[secret_name] = secret
-        binary_output = json.dumps(secrets).encode()
-        tarinfo = tarfile.TarInfo(name=filename)
-        tarinfo.size = len(binary_output)
-        file_dumps.append((binary_output, tarinfo))
-    with tarfile.open("/tmp/backup.tar.gz", "w:gz") as tar:
-        print(f"Writing files to /tmp/backup.tar.gz...")
-        for binary_output, tarinfo in file_dumps:
-            tar.addfile(tarinfo, fileobj=io.BytesIO(binary_output))
-    size = get_formatted_size("/tmp/backup.tar.gz")
-    print(f"Backup has been written to /tmp/backup.tar.gz with a size of {size}")
+    rmtree(DOWNLOAD_TMP_PATH, ignore_errors=True)
+    os.makedirs(DOWNLOAD_TMP_PATH)
+    restore_dir = os.path.join(DOWNLOAD_TMP_PATH, "restore")
+    try:
+        if path:
+            path = path.removesuffix("/")
+            files = list_files(snapshot, path)
+            filetype = [f.get("type") for f in files if f.get("path") == path][0]
+            filename = Path(path).name
+            print(f"Restoring {path} from snapshot '{snapshot}'")
+            restic_restore(snapshot_id=snapshot, include=[path], target_dir=restore_dir)
+            if filetype == "dir":
+                filename = filename + ".tar"
+                create_tar(restore_dir, path, os.path.join(DOWNLOAD_TMP_PATH, filename))
+            else:
+                os.replace(
+                    f"{restore_dir}{path}", os.path.join(DOWNLOAD_TMP_PATH, filename)
+                )
+            rmtree(restore_dir)
+            file_dumps.append(filename)
+        if volumes:
+            if SERVICE == "ALL":
+                logger.error("Please specify '--host' when using '--volumes'")
+                exit(1)
+            files = list_files(snapshot, VOLUME_PATH)
+            for f in files[1:]:
+                path = f["path"]
+                if Path(path).name.startswith(SERVICE) and f["type"] == "dir":
+                    filename = f"{Path(path).name}.tar"
+                    print(f"Restoring {path} from snapshot '{snapshot}'")
+                    restic_restore(snapshot_id=snapshot, include=[path], target_dir=restore_dir)
+                    create_tar(restore_dir, path, os.path.join(DOWNLOAD_TMP_PATH, filename))
+                    rmtree(restore_dir)
+                    file_dumps.append(filename)
+        if secrets:
+            if SERVICE == "ALL":
+                logger.error("Please specify '--host' when using '--secrets'")
+                exit(1)
+            filename = f"{SERVICE}.json"
+            print(f"Restoring {SECRET_PATH} from snapshot '{snapshot}'")
+            restic_restore(
+                snapshot_id=snapshot,
+                include=[SECRET_PATH.removesuffix("/")],
+                target_dir=restore_dir,
+            )
+            secrets = read_secrets(f"{restore_dir}{SECRET_PATH}")
+            rmtree(restore_dir)
+            with open(os.path.join(DOWNLOAD_TMP_PATH, filename), "w") as file:
+                json.dump(secrets, file)
+            file_dumps.append(filename)
+        print(f"Writing files to {BACKUP_FILE}...")
+        with tarfile.open(BACKUP_FILE, "w:gz") as tar:
+            for filename in file_dumps:
+                tar.add(os.path.join(DOWNLOAD_TMP_PATH, filename), arcname=filename)
+    finally:
+        rmtree(DOWNLOAD_TMP_PATH, ignore_errors=True)
+    size = get_formatted_size(BACKUP_FILE)
+    print(f"Backup has been written to {BACKUP_FILE} with a size of {size}")
+
+
+def create_tar(restore_dir, path, tar_file):
+    # The archive contains the children of the given directory,
+    # named by their full path without the leading slash.
+    with tarfile.open(tar_file, "w") as tar:
+        for child in sorted(os.listdir(f"{restore_dir}{path}")):
+            tar.add(
+                f"{restore_dir}{path}/{child}",
+                arcname=f"{path}/{child}".removeprefix("/"),
+            )
+
+
+def read_secrets(secret_dir):
+    secrets = {}
+    for name in sorted(os.listdir(secret_dir)):
+        secret_path = os.path.join(secret_dir, name)
+        if name.startswith(SERVICE) and os.path.isfile(secret_path):
+            with open(secret_path) as file:
+                secrets[name.removeprefix(f"{SERVICE}_")] = file.read()
+    return secrets
 
 
 def get_formatted_size(file_path):
@@ -632,18 +665,6 @@ def get_formatted_size(file_path):
     return f"{round(file_size, 3)} {units[-1]}"
 
 
-def dump(snapshot, path):
-    cmd = restic.cat.base_command() + ["dump"]
-    cmd = cmd + [snapshot, path]
-    print(f"Dumping {path} from snapshot '{snapshot}'")
-    output = subprocess.run(cmd, capture_output=True)
-    if output.returncode:
-        logger.error(
-            f"error while dumping {path} from snapshot '{snapshot}': {output.stderr}"
-        )
-        exit(1)
-    return output.stdout
-
 
 if __name__ == "__main__":
     cli()