also allow running the update script locally

1 files changed, 114 insertions, 0 deletions

diff --git a/bin/for-servers b/bin/for-servers
new file mode 100755
index 0000000..4a975cf
--- /dev/null
+++ b/bin/for-servers
@@ -0,0 +1,114 @@
+#!/usr/bin/env python3
+import asyncio
+import os
+import sys
+from typing import Any, Callable, Coroutine
+
+
+HOST_COLORS = {
+    "tozt": "\x1b[1;33m",
+    "mail": "\x1b[32m",
+    "partofme": "\x1b[35m",
+}
+
+
+def color_host(host: str):
+    return HOST_COLORS[host] + host + "\x1b[m"
+
+
+async def unlock_rbw():
+    proc = await asyncio.create_subprocess_exec("rbw", "unlock")
+    await proc.wait()
+
+
+async def get_password(host: str):
+    proc = await asyncio.create_subprocess_exec(
+        "rbw",
+        "get",
+        host,
+        os.environ["USER"],
+        stdout=asyncio.subprocess.PIPE,
+    )
+
+    stdout, _ = await proc.communicate()
+    return stdout.rstrip()
+
+
+async def read_stream(
+    stream: asyncio.StreamReader,
+    print_cb: Callable[[str], None],
+    sudo_cb: Coroutine[Any, Any, None] | None,
+):
+    buf = b""
+    while True:
+        read = await stream.read(1024)
+        if len(read) == 0:
+            if len(buf) > 0:
+                print_cb(buf.decode())
+            break
+
+        buf += read
+        lines = buf.split(b"\n")
+        buf = lines.pop()
+        for line in lines:
+            print_cb(line.decode())
+
+        if sudo_cb and buf == f"[sudo] password for {os.environ['USER']}: ".encode():
+            await sudo_cb
+            sudo_cb = None
+            buf = b""
+
+
+async def puppet_tozt(host: str):
+    password = await get_password(host)
+
+    proc = await asyncio.create_subprocess_exec(
+        "ssh",
+        host,
+        "sudo",
+        "--stdin",
+        *sys.argv[1:],
+        stdin=asyncio.subprocess.PIPE,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE,
+    )
+    assert proc.stdout is not None
+    assert proc.stderr is not None
+
+    colored_host = color_host(host)
+
+    async def sudo_cb():
+        assert proc.stdin is not None
+        proc.stdin.write(password + b"\n")
+        await proc.stdin.drain()
+
+    await asyncio.gather(
+        read_stream(
+            proc.stdout,
+            lambda line: print(f"[{colored_host}:out] {line}"),
+            None,
+        ),
+        read_stream(
+            proc.stderr,
+            lambda line: print(f"[{colored_host}:\x1b[31merr\x1b[m] {line}"),
+            sudo_cb(),
+        ),
+    )
+
+    ret = await proc.wait()
+    if ret == 0:
+        print(f"[{colored_host}] Exited successfully")
+    else:
+        print(f"[{colored_host}] \x1b[31mExited with code {ret}\x1b[m")
+
+
+async def main():
+    await unlock_rbw()
+    await asyncio.gather(
+        puppet_tozt("tozt"),
+        puppet_tozt("mail"),
+        puppet_tozt("partofme"),
+    )
+
+
+asyncio.run(main())