also allow running the update script locally

3 files changed, 124 insertions, 111 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())
diff --git a/bin/puppet-tozt b/bin/puppet-tozt
index 0ce88a7..bdf0460 100755
--- a/bin/puppet-tozt
+++ b/bin/puppet-tozt
@@ -1,113 +1,6 @@
-#!/usr/bin/env python3
-import asyncio
-import os
-from typing import Any, Callable, Coroutine
+#!/bin/sh
+set -eu
 
+cd "$(dirname "$0")/.."
 
-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",
-        "puppet-tozt",
-        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())
+exec bin/for-servers puppet-tozt
diff --git a/bin/update b/bin/update
new file mode 100755
index 0000000..0784346
--- /dev/null
+++ b/bin/update
@@ -0,0 +1,6 @@
+#!/bin/sh
+set -eu
+
+cd "$(dirname "$0")/.."
+
+exec bin/for-servers update