add script to puppet all machines at once

1 files changed, 93 insertions, 0 deletions

diff --git a/bin/puppet-tozt b/bin/puppet-tozt
new file mode 100755
index 0000000..35579b3
--- /dev/null
+++ b/bin/puppet-tozt
@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+import asyncio
+import os
+from typing import Any, Callable, Coroutine
+
+
+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 == b"[sudo] password for doy: ":
+            await sudo_cb
+            sudo_cb = None
+
+
+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
+
+    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"[{host}:out] {line}"),
+            None,
+        ),
+        read_stream(
+            proc.stderr,
+            lambda line: print(f"[{host}:err] {line}"),
+            sudo_cb(),
+        ),
+    )
+
+    ret = await proc.wait()
+    if ret == 0:
+        print(f"[{host}] Exited successfully")
+    else:
+        print(f"[{host}] Exited with code {ret}")
+
+
+async def main():
+    await asyncio.gather(
+        puppet_tozt("tozt"),
+        puppet_tozt("mail"),
+        puppet_tozt("partofme"),
+    )
+
+
+asyncio.run(main())