#!/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())
