From bf0d278243acbcdcb4f1447541c3248212a02c0d Mon Sep 17 00:00:00 2001 From: Jesse Luehrs Date: Sun, 3 Mar 2024 09:57:19 -0500 Subject: also allow running the update script locally --- bin/for-servers | 114 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ bin/puppet-tozt | 115 ++------------------------------------------------------ bin/update | 6 +++ 3 files changed, 124 insertions(+), 111 deletions(-) create mode 100755 bin/for-servers create mode 100755 bin/update 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 -- cgit v1.2.3-54-g00ecf