| 1 | #!/usr/bin/env python3 |
| 2 | import argparse |
| 3 | import hashlib |
| 4 | import json |
| 5 | import os |
| 6 | from pathlib import Path |
| 7 | import subprocess |
| 8 | import time |
| 9 | |
| 10 | import router |
| 11 | |
| 12 | |
| 13 | def main(): |
| 14 | parser = argparse.ArgumentParser() |
| 15 | parser.add_argument("--ready-file", type=Path, required=True) |
| 16 | parser.add_argument("--timeout", type=int, default=180) |
| 17 | args = parser.parse_args() |
| 18 | assert 1 <= args.timeout <= 600 |
| 19 | args.ready_file.unlink(missing_ok=True) |
| 20 | stop = args.ready_file.with_suffix(".stop") |
| 21 | stop.unlink(missing_ok=True) |
| 22 | routes = Path(router.ROUTES) |
| 23 | original = routes.read_bytes() |
| 24 | metadata = routes.stat() |
| 25 | candidate = router.render(Path(router.TOKEN).read_text().strip()) |
| 26 | start = "globe." + os.environ["STUDIO_DOMAIN"] + " {" |
| 27 | begin = original.index(start.encode()) |
| 28 | depth = 0 |
| 29 | end = None |
| 30 | for position in range(begin + len(start) - 1, len(original)): |
| 31 | depth += (original[position] == ord("{")) - (original[position] == ord("}")) |
| 32 | if depth == 0: |
| 33 | end = position + 1 |
| 34 | break |
| 35 | assert end is not None |
| 36 | replacement = candidate[candidate.index(start):candidate.index("dashboard.internal." + os.environ["STUDIO_DOMAIN"] + ":")].strip().encode() |
| 37 | content = original[:begin] + replacement + original[end:] |
| 38 | active = subprocess.run(["systemctl", "is-active", "studio-router"], capture_output=True).returncode == 0 |
| 39 | pending = routes.with_name("routes.browser-test.pending") |
| 40 | |
| 41 | def install(value): |
| 42 | pending.write_bytes(value) |
| 43 | os.chown(pending, metadata.st_uid, metadata.st_gid) |
| 44 | pending.chmod(metadata.st_mode & 0o777) |
| 45 | pending.replace(routes) |
| 46 | subprocess.run(["systemctl", "reload", "caddy"], check=True, capture_output=True, timeout=30) |
| 47 | |
| 48 | try: |
| 49 | if active: |
| 50 | subprocess.run(["systemctl", "stop", "studio-router"], check=True, capture_output=True, timeout=30) |
| 51 | install(content) |
| 52 | args.ready_file.write_text(json.dumps({"origin": "https://globe." + os.environ["STUDIO_DOMAIN"], "ready": True}) + "\n") |
| 53 | args.ready_file.chmod(0o600) |
| 54 | deadline = time.monotonic() + args.timeout |
| 55 | while not stop.exists() and time.monotonic() < deadline: |
| 56 | time.sleep(.2) |
| 57 | finally: |
| 58 | try: |
| 59 | install(original) |
| 60 | assert hashlib.sha256(routes.read_bytes()).digest() == hashlib.sha256(original).digest() |
| 61 | finally: |
| 62 | if active: |
| 63 | subprocess.run(["systemctl", "start", "studio-router"], check=True, capture_output=True, timeout=30) |
| 64 | pending.unlink(missing_ok=True) |
| 65 | args.ready_file.unlink(missing_ok=True) |
| 66 | stop.unlink(missing_ok=True) |
| 67 | print(json.dumps({"original_public_routes_restored": True, "router_state_restored": True})) |
| 68 | |
| 69 | |
| 70 | if __name__ == "__main__": |
| 71 | main() |