1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
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())
|