1#!/usr/bin/env python3
2import argparse
3import json
4import os
5from pathlib import Path
6import sqlite3
7import subprocess
8import time
9import urllib.error
10import urllib.request
11
12
13parser = argparse.ArgumentParser(description="Copy retained app files on the new host before starting the service.")
14parser.add_argument("service", choices=["tailscale", "ddns-updater", "copyparty", "ytdl", "pds"])
15
16
17def main():
18 service = parser.parse_args().service
19 if os.geteuid() != 0:
20 raise SystemExit("Run this importer as root on the production host.")
21 token = Path("/var/lib/studio/nomad.token").read_text().strip()
22 request = urllib.request.Request(
23 "http://127.0.0.1:4646/v1/job/" + service,
24 headers={"X-Nomad-Token": token},
25 )
26 try:
27 with urllib.request.urlopen(request, timeout=10) as response:
28 job = json.load(response)
29 if not job.get("Stop"):
30 raise SystemExit("Stop " + service + " before copying its files.")
31 except urllib.error.HTTPError as error:
32 if error.code != 404:
33 raise
34 request.full_url += "/allocations"
35 try:
36 with urllib.request.urlopen(request, timeout=10) as response:
37 allocations = json.load(response)
38 except urllib.error.HTTPError as error:
39 if error.code != 404:
40 raise
41 allocations = []
42 if any(item["ClientStatus"] not in ("complete", "failed", "lost") for item in allocations):
43 raise SystemExit("Wait for " + service + " allocations to stop before copying its files.")
44
45 root = Path("/srv/prod") / service
46 dataset = subprocess.check_output(["findmnt", "-n", "-o", "SOURCE", "--mountpoint", str(root)], text=True).strip()
47 if dataset != "globe/prod/" + service:
48 raise SystemExit("Mount globe/prod/" + service + " at " + str(root) + " before copying files.")
49 uid = 0 if service == "tailscale" else json.loads(Path("/var/lib/studio/identities.json").read_text())[service]
50 directories = {
51 "tailscale": ("tailscale", "var/lib/tailscale"),
52 "ddns-updater": ("ddns-updater", "updater/data"),
53 "copyparty": ("copyparty/copyparty", "cfg/copyparty"),
54 "ytdl": ("ytdl-sub", "config"),
55 "pds": ("pds", "pds"),
56 }
57 source_name, target_name = directories[service]
58 source = Path("/mnt/storage1/apps") / source_name
59 target = root / target_name
60 if not source.is_dir():
61 raise SystemExit("Retained app files are missing at " + str(source))
62 if service == "tailscale":
63 json.loads((source / "tailscaled.state").read_text())
64 elif service == "ddns-updater":
65 json.loads((source / "updates.json").read_text())
66 elif service == "copyparty":
67 for database in source.glob("*.db"):
68 with sqlite3.connect("file:" + str(database) + "?mode=ro", uri=True) as connection:
69 if connection.execute("PRAGMA integrity_check").fetchone()[0] != "ok":
70 raise SystemExit("Check the retained Copyparty database before copying " + database.name)
71 if not (source / "shares.db").is_file():
72 raise SystemExit("Restore the retained Copyparty shares.db before importing.")
73 elif service == "pds":
74 for name in ("account", "sequencer", "did_cache"):
75 if not (source / (name + ".sqlite")).is_file():
76 raise SystemExit("Restore the retained PDS " + name + " database before importing.")
77 if list(source.rglob("*-wal")) or list(source.rglob("*-shm")):
78 raise SystemExit("Back up all PDS SQLite databases, including actor stores, before importing WAL files.")
79 for database in source.rglob("*.sqlite"):
80 with sqlite3.connect("file:" + str(database) + "?mode=ro&immutable=1", uri=True) as connection:
81 if connection.execute("PRAGMA integrity_check").fetchone()[0] != "ok":
82 raise SystemExit("Check the retained PDS database before importing " + database.name)
83
84 snapshot = dataset + "@before-files-import-" + str(time.time_ns())
85 subprocess.run(["zfs", "snapshot", snapshot], check=True)
86 target.mkdir(parents=True, exist_ok=True)
87 exclusions = ["--exclude=*.lock", "--exclude=/.ytdl-sub-lock", "--exclude=/.cache/", "--exclude=/work/"]
88 subprocess.run(["rsync", "-aH", "--delete", *exclusions, str(source) + "/", str(target) + "/"], check=True)
89 drift = subprocess.check_output([
90 "rsync", "-rlnc", "--delete", "--out-format=%n", *exclusions, str(source) + "/", str(target) + "/",
91 ], text=True)
92 if drift:
93 raise SystemExit("Copied files differ from the retained source. Restore " + snapshot + " before retrying.")
94 os.chown(target, uid, uid)
95 for directory, children, files in os.walk(target):
96 for name in children + files:
97 os.chown(Path(directory) / name, uid, uid, follow_symlinks=False)
98 if service == "tailscale":
99 os.chmod(target / "tailscaled.state", 0o600)
100 elif service == "copyparty":
101 for database in target.glob("*.db"):
102 with sqlite3.connect("file:" + str(database) + "?mode=ro", uri=True) as connection:
103 if connection.execute("PRAGMA integrity_check").fetchone()[0] != "ok":
104 raise SystemExit("Copied Copyparty database needs recovery from " + snapshot)
105 elif service == "pds":
106 for database in target.rglob("*.sqlite"):
107 with sqlite3.connect("file:" + str(database) + "?mode=ro&immutable=1", uri=True) as connection:
108 if connection.execute("PRAGMA integrity_check").fetchone()[0] != "ok":
109 raise SystemExit("Copied PDS database needs recovery from " + snapshot)
110 print(json.dumps({"service": service, "snapshot": snapshot, "copied": True}))
111
112
113if __name__ == "__main__":
114 main()