also allow running the update script locally

1 files changed, 4 insertions, 111 deletions

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