1#!/usr/bin/env python3
2"""
3Wrap the devcontrol binary used for running a local instance of the
4Tailscale control plane.
5
6I use a tiny subset of the available flags and scenarios; this script
7wraps them in a friendly interface.
8"""
9
10import argparse
11import atexit
12import os
13from pathlib import Path
14import subprocess
15import sys
16from typing import TypedDict
17
18from chives.text import coloured
19
20
21Args = TypedDict(
22    "args",
23    {
24        "audit_logs": bool,
25        "persist": bool,
26        "tfacc": bool,
27    },
28)
29
30
31def parse_args():
32    parser = argparse.ArgumentParser(
33        prog="devcontrol",
34        description="Run devcontrol with different options",
35    )
36
37    parser.add_argument("--audit-logs", action="store_true", help="save audit logs")
38    # TODO: Can I make this a flag to store state in a specific dir?
39    parser.add_argument(
40        "--persist",
41        action="store_true",
42        help="save the devcontrol state in /tmp/devcontrol",
43    )
44    parser.add_argument(
45        "--tfacc",
46        action="store_true",
47        help="run the Terraform acceptance test scenario",
48    )
49
50    args = parser.parse_args()
51
52    return {
53        "audit_logs": args.audit_logs,
54        "persist": args.persist,
55        "tfacc": args.tfacc,
56    }
57
58
59if __name__ == "__main__":
60    args = parse_args()
61
62    corp_dir = Path.home() / "repos/corp"
63
64    tailcontrol_cmd = ["./tool/go", "run", "--tags=tailscale_saas", "./cmd/devcontrol"]
65    tailcontrol_env = {
66        "HOME": str(Path.home()),
67        "PATH": os.environ["PATH"],
68    }
69
70    # See https://www.notion.so/tailscale/Running-a-local-CONTROL-clients-and-DERP-d69826480c704c1f9f39950478e7303f#f744c21f49f14c4fa748bff0a944e9bb
71    if args["audit_logs"]:
72        # Compile logzservice before running it, so the terminate() signal
73        # goes to the logzservice process rather than `go run`.
74        #
75        # Otherwise, the compiled logzservice binary can be left running
76        # after this script exits.
77        subprocess.check_call(["./tool/go", "build", "./cmd/logzservice"], cwd=corp_dir)
78        logz_proc = subprocess.Popen(["./logzservice", "-dev"], cwd=corp_dir)
79        atexit.register(lambda: logz_proc.terminate())
80        tailcontrol_env.update(
81            {
82                "TS_LOG_TARGET": "http://localhost:9951",
83                "TS_DEBUG_CONTROL_LOGTAIL_UPLOAD": "true",
84            }
85        )
86
87        def stop_logz_proc() -> None:
88            logz_proc.terminate()
89
90        atexit.register(stop_logz_proc)
91
92    if args["persist"]:
93        tailcontrol_cmd.append("--dir=/tmp/devcontrol")
94
95    if args["tfacc"]:
96        tailcontrol_cmd.append("--generate-test-devices=terraform-acceptance-testing")
97
98    debug_string = ["->"]
99    for name, value in tailcontrol_env.items():
100        if name in {"HOME", "PATH"}:
101            continue
102        debug_string.append(f"{name}={value}")
103    debug_string.extend(tailcontrol_cmd)
104
105    print(coloured(" ".join(debug_string), "blue"))
106    try:
107        subprocess.check_call(tailcontrol_cmd, env=tailcontrol_env, cwd=corp_dir)
108    except KeyboardInterrupt:
109        sys.exit(0)