-
Notifications
You must be signed in to change notification settings - Fork 0
/
controller.py
154 lines (136 loc) · 4.85 KB
/
controller.py
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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
import os, sys
def get_ip():
result = os.popen("ifconfig | grep addr:10")
ips = result.read().strip().split()
for each in ips:
if each[:5] == "addr:":
return 0, each[5:]
return 0, "0.0.0.0"
def update_with_screen(lines=None):
if lines is None:
with open("servers.csv", "r")as fin:
lines = fin.readlines()
for each in lines[1:]:
data = each.split(",")
with open("run.sh", "w") as fout:
fout.write(
f"""#!/usr/bin/expect
set host {data[0]}
set user {data[1]}
set timeout 2
spawn ssh $user@$host
expect "*password:*"
send "{data[2]}\r"
expect "*]#"
send "screen -S nvidia_reporter -X quit\r"
expect "*]#"
send "rm -rf .nvidia_reporter{data[0][-3:]}/\r"
expect "*]#"
send "mkdir .nvidia_reporter\r"
expect "*]#"
send "cd .nvidia_reporter\r"
expect "*]#"
send "wget http://{base_url}/reporter\r"
expect "*]#"
send "mv reporter reporter.py\r"
expect "*]#"
send "screen -S nvidia_reporter\r"
expect "*]#"
send "python reporter.py {base_url}\r"
expect "*]#"
send "exit\r"
expect eof""")
os.system("chmod 777 run.sh && ./run.sh")
print(f"\n====={data[0]}, ok!=====\n")
os.system("rm run.sh")
def stop_screen():
with open("servers.csv", "r")as fin:
lines = fin.readlines()
for each in lines[1:]:
data = each.strip('\n').split(",")
with open("run.sh", "w") as fout:
fout.write(
f"""#!/usr/bin/expect
set host {data[0]}
set user {data[1]}
set timeout 2
spawn ssh $user@$host
expect "*password:*"
send "{data[2]}\r"
expect "*]#"
send "screen -S nvidia_reporter -X quit\r"
expect "*]#"
send "rm -rf .nvidia_reporter{data[0][-3:]}/\r"
expect "*]#"
send "exit\r"
expect eof""")
os.system("chmod 777 run.sh && ./run.sh")
print(f"\n====={data[0]}, ok!=====\n")
os.system("rm run.sh")
def create_users():
shell_context = """#!/bin/sh
echo \\\"Remove old account if exist\\\"
userdel -r -f \\\"\\$1\\\" || true
echo \\\"Add new account for \\$1\\\"
useradd -m \\$1 -d /remote-home/\\$1 -u \\$2 -U
echo \\\"Change password for \\$1\\\"
echo \\\"\\$1\\n\\$1\\\" | passwd \\\"\\$1\\\"
mkdir /remote-home/\\$1
chown \\$1:\\$1 /remote-home/\\$1
echo \\\"Add \\$1 to group sudo\\\"
chsh \\$1 -s /bin/bash
echo \\\"DONE\\\" """
with open("users.csv", "r")as fin:
lines = fin.readlines()
commands = ""
for line in lines[1:]:
if line.startswith("#"):
continue
line = line.strip('\n').split(',')
commands += f"""send "sudo bash ./add_remote.sh {line[0]} {line[1]}\n" \nexpect "*]#"\n"""
with open("servers.csv", "r")as fin:
lines = fin.readlines()
for each in lines[1:]:
data = each.strip('\n').split(",")
with open("run.sh", "w") as fout:
fout.write(
f"""#!/usr/bin/expect
set host {data[0]}
set user {data[1]}
set timeout 2
spawn ssh $user@$host
expect "*password:*"
send "{data[2]}\r"
expect "*]#"
send "sudo su\r"
expect "*password:*"
send "{data[2]}\r"
expect "*]#"
send "echo -e '"""
+ shell_context +
""" '> add_remote.sh\r"
expect "*]#"\n"""
+ commands +
"""send "rm add_remote.sh\r"
expect "*]#"
send "exit\r"
expect eof""")
os.system("chmod 700 run.sh && ./run.sh")
print(f"\n====={data[0]}, ok!=====\n")
os.system("rm run.sh")
if __name__ == "__main__":
err, ip = get_ip()
if err > 0:
print("Can't get the IP address.")
else:
base_url = ip + ':8997'
if len(sys.argv) == 1:
print("-u Update all reporters (with servers.csv)")
print("-s Stop all reporters (with servers.csv)")
print("-c Create remote-home for users (with users.csv)")
elif sys.argv[1] == '-u':
update_with_screen()
elif sys.argv[1] == '-s':
stop_screen()
elif sys.argv[1] == '-c':
create_users()