-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmockdevice.py
181 lines (138 loc) · 6.13 KB
/
mockdevice.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
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
#!/usr/bin/python
#
import sys, time, os
#import psutil
import MockSSH
import MockSSHExtensions
import traceback
from random import randint
from time import sleep
import json
from collections import defaultdict
from twisted.python import log
from threading import Thread
from twisted.internet import reactor
from twisted.internet import reactor
from twisted.internet.protocol import ClientFactory, ServerFactory
from twisted.conch.telnet import TelnetTransport, TelnetProtocol, AuthenticatingTelnetProtocol, ITelnetProtocol, TelnetBootstrapProtocol, StatefulTelnetProtocol
from twisted.cred import checkers, portal
from twisted.conch import avatar, interfaces as conchinterfaces, recvline
from twisted.conch import manhole
from twisted.conch.insults import insults
from twisted.conch.ssh import session
import multiprocessing
import subprocess
pool = None
batch_size = 100 # ports to handle per process
def main():
# FIXME Make proper arguments
port_low = int(sys.argv[2])
port_high = int(sys.argv[3])
all_ports = port_high - port_low
last_batch = all_ports % batch_size
processes = all_ports / batch_size
protocol_type = sys.argv[4]
device_file = sys.argv[5]
interface = sys.argv[1]
log.startLogging(sys.stdout)
with open(device_file) as f:
data = json.load(f)
global pool
process_pool_size = processes if last_batch <=0 else (processes + 1)
if process_pool_size == 1:
print "Running in a single process"
spawn_server(port_low, port_high, interface, protocol_type, data)
else:
print "Spawning pool with %s processes" % process_pool_size
pool = multiprocessing.Pool(processes if last_batch <=0 else (processes + 1))
args = []
for i in range(0, processes):
starting_port = port_low + i*batch_size
ending_port = port_low + batch_size + i*batch_size
print "Spawning process for ports: %s - %s" % (starting_port, ending_port)
if i == processes - 1:
ending_port = ending_port + 1
args.append((i, starting_port, ending_port, interface, protocol_type, data))
r = pool.map_async(spawn_server_wrapper, args)
if last_batch > 0:
starting_port = port_low + processes*batch_size + 1
ending_port = port_high + 1
print "Spawning process for ports: %s - %s" % (starting_port, ending_port)
r2 = pool.map_async(spawn_server_wrapper, [(processes, starting_port, ending_port, interface, protocol_type, data)])
try:
r2.wait()
except KeyboardInterrupt:
r.wait()
r.wait()
def spawn_server_wrapper(args):
# Wait N * 500 milliseconds to prevent all sub processes to start at once (e.g. batch 7 would wait 3.5 seconds before starting)
# If processes start at once, this error might occur: https://twistedmatrix.com/trac/ticket/4759
sleep((500.0 * args[0]) / 1000)
spawn_server(*args[1:])
def spawn_server(port_low, port_high, interface, protocol_type, data):
#p = psutil.Process(os.getpid())
try:
os.nice(-1)
except Exception as e:
print "Unable to lower niceness(priority), RUN AS SUDO, running with normal priority"
for i in range(port_low, port_high):
try:
print "Spawning "
show_commands, prompt_change_commands, usr, passwd, cmd_delay, default_prompt = parse_commands(data)
users = {usr : passwd}
local_commands = []
for cmd in show_commands:
command = getShowCommand(cmd, data, show_commands[cmd], cmd_delay)
local_commands.append(command)
for cmd in prompt_change_commands:
if ("password" in prompt_change_commands[cmd]):
command = getPasswordPromptCommand(cmd, prompt_change_commands[cmd], cmd_delay)
else:
command = getPromptChangingCommand(cmd, prompt_change_commands[cmd], cmd_delay)
local_commands.append(command)
factory = None
if (protocol_type == "ssh"):
factory = MockSSH.getSSHFactory(local_commands, default_prompt, ".", **users)
elif (protocol_type == "telnet"):
factory = MockSSHExtensions.getTelnetFactory(local_commands, default_prompt, **users)
reactor.listenTCP(i, factory, interface=interface)
except Exception as e:
print >> sys.stderr, traceback.format_exc()
print "Unable to open port at %s, due to: %s" % (i, e)
reactor.run()
def parse_commands(data):
show_commands = defaultdict(list)
prompt_change_commands = {}
default_prompt = data["setting_default_prompt"]
cmd_delay = data["setting_cmd_delay"]
usr = data["setting_default_user"]
passwd = data["setting_default_passwd"]
print "Using username: %s and password: %s" % (usr, passwd)
for cmd in data:
if isinstance(data[cmd], dict):
prompt_change_commands[cmd] = data[cmd]
#print "Adding prompt changing command: %s" % cmd
else:
cmd_split = cmd.split(" ", 1)
if (len(cmd_split) == 1):
show_commands[cmd_split[0]]
else:
show_commands[cmd_split[0]].append(cmd_split[1])
#print "Adding show command: %s, with arguments: %s" % (cmd, show_commands[cmd])
return (show_commands, prompt_change_commands, usr, passwd, cmd_delay, default_prompt)
def getPasswordPromptCommand(cmd, values, cmd_delay):
return MockSSHExtensions.SimplePromptingCommand(values["name"],values["password"],values["prompt"],values["newprompt"],values["error_message"], cmd_delay)
def getPromptChangingCommand(cmd, values, cmd_delay):
return MockSSHExtensions.PromptChangingCommand(cmd, values["newprompt"], cmd_delay)
def getShowCommand(cmd, data, arguments, cmd_delay):
return MockSSHExtensions.ShowCommand(cmd, data, cmd_delay, *arguments)
if __name__ == "__main__":
try:
main()
except KeyboardInterrupt:
print "User interrupted"
global pool
pool.close()
pool.terminate()
pool.join()
sys.exit(1)