class TestTcpServer(object):
def __init__(self, server_addr):
self._tcp_server = ZTcpServer(('127.0.0.1', 9876))
self._codec = ZLengthHeadCodec(self._on_string_message)
self._guard = Condition(Lock())
self._connections = {}
self._tcp_server.set_connection_callback(self._on_connection)
self._tcp_server.set_message_callback(self._codec.on_message)
def start(self):
self._tcp_server.start()
self._tcp_server.serve_forever()
def set_thread_num(self, num):
self._tcp_server.set_thread_num(num)
def _on_connection(self, tcp_conn):
logging.info(tcp_conn.name() + ' is ' + (tcp_conn.connected() and 'UP' or 'DOWN'))
with self._guard:
if tcp_conn.connected():
self._connections[tcp_conn.name()] = tcp_conn
else:
self._connections.pop(tcp_conn.name())
def _on_string_message(self, tcp_conn, message, receive_time):
with self._guard:
for conn_name in self._connections:
self._codec.send(self._connections[conn_name], message)
def __init__(self, server_addr):
self._tcp_server = ZTcpServer(('127.0.0.1', 9876))
self._codec = ZLengthHeadCodec(self._on_string_message)
self._guard = Condition(Lock())
self._connections = {}
self._tcp_server.set_connection_callback(self._on_connection)
self._tcp_server.set_message_callback(self._codec.on_message)
class Scanner(object):
def __init__(self, event_loop, sid, server_addr):
self._sid = sid
self._codec = ZLengthHeadCodec(self._on_string_message)
self._tcp_client = ZTcpClient(event_loop, server_addr)
self._tcp_client.set_connection_callback(self._on_connection)
self._tcp_client.set_connection_error_callback(self._on_error_connection)
self._tcp_client.set_message_callback(self._codec.on_message)
self._guard = Condition(Lock())
self._connection_callback = None
self._connection_error_callback = None
self._message_callback = None
def connect(self):
self._tcp_client.connect()
def disconnect(self):
self._tcp_client.disconnect()
def set_connection_callback(self, cb):
self._connection_callback = cb
def set_connection_error_callback(self, cb):
self._connection_error_callback = cb
def set_message_callback(self, cb):
self._message_callback = cb
def send(self, tcp_conn, msg):
self._codec.send(tcp_conn, msg)
def get_id(self):
return self._sid
def _on_connection(self, tcp_conn):
if self._connection_callback:
self._connection_callback(self, tcp_conn)
def _on_error_connection(self, errno):
if self._connection_error_callback:
self._connection_error_callback(self, errno)
else:
logging.error('failed to connect to %s, errno=(%d)' %str(server_addr), errno)
def _on_string_message(self, tcp_conn, message, receive_time):
if self._message_callback:
self._message_callback(self, tcp_conn, message, receive_time)
def __init__(self, server_addr):
self._event_loop_thread = ZEventLoopThread()
self._tcp_client = ZTcpClient(self._event_loop_thread.start_loop(), server_addr)
self._codec = ZLengthHeadCodec(self._on_string_message)
self._tcp_client.set_connection_callback(self._on_connection)
self._tcp_client.set_message_callback(self._codec.on_message)
self._tcp_conn = None
self._guard = Condition(Lock())
def __init__(self, server_addr, cfg_file, scan_list_dir):
self._cfg_parser = ConfigParser.ConfigParser()
self._cfg_parser.read(cfg_file)
self._construct_scan_parameters()
self._scan_list_dir = scan_list_dir
self._tcp_server = ZTcpServer(server_addr)
self._codec = ZLengthHeadCodec(self._on_string_message)
self._guard = Condition(Lock())
self._tcp_server.set_connection_callback(self._on_connection)
self._tcp_server.set_message_callback(self._codec.on_message)
self._hostname = gethostname()
def __init__(self, event_loop, sid, server_addr):
self._sid = sid
self._codec = ZLengthHeadCodec(self._on_string_message)
self._tcp_client = ZTcpClient(event_loop, server_addr)
self._tcp_client.set_connection_callback(self._on_connection)
self._tcp_client.set_connection_error_callback(self._on_error_connection)
self._tcp_client.set_message_callback(self._codec.on_message)
self._guard = Condition(Lock())
self._connection_callback = None
self._connection_error_callback = None
self._message_callback = None
class TestTcpClient(object):
def __init__(self, server_addr):
self._event_loop_thread = ZEventLoopThread()
self._tcp_client = ZTcpClient(self._event_loop_thread.start_loop(), server_addr)
self._codec = ZLengthHeadCodec(self._on_string_message)
self._tcp_client.set_connection_callback(self._on_connection)
self._tcp_client.set_message_callback(self._codec.on_message)
self._tcp_conn = None
self._guard = Condition(Lock())
def connect(self):
self._tcp_client.connect()
def send_message_until_quit(self):
while True:
try:
s = input()
if s.find("bye") < 0:
self._write(s)
else:
self._tcp_client.disconnect()
return
except Exception as e:
logging.error("input error %s" % (str(e)))
def _write(self, s):
with self._guard:
if self._tcp_conn is not None:
self._codec.send(self._tcp_conn, s)
def _on_connection(self, tcp_conn):
logging.info(tcp_conn.name() + " is " + (tcp_conn.connected() and "UP" or "DOWN"))
with self._guard:
if tcp_conn.connected():
self._tcp_conn = tcp_conn
else:
self._tcp_conn = None
def _on_string_message(self, tcp_conn, message, recieve_time):
logging.info(message)
class Scanner(object):
def __init__(self, server_addr, cfg_file, scan_list_dir):
self._cfg_parser = ConfigParser.ConfigParser()
self._cfg_parser.read(cfg_file)
self._construct_scan_parameters()
self._scan_list_dir = scan_list_dir
self._tcp_server = ZTcpServer(server_addr)
self._codec = ZLengthHeadCodec(self._on_string_message)
self._guard = Condition(Lock())
self._tcp_server.set_connection_callback(self._on_connection)
self._tcp_server.set_message_callback(self._codec.on_message)
self._hostname = gethostname()
def start(self):
self._tcp_server.start()
self._tcp_server.serve_forever()
def set_thread_num(self, num):
self._tcp_server.set_thread_num(num)
def _on_connection(self, tcp_conn):
logging.info('Scanner (%s) is %s' %(tcp_conn.name(), (tcp_conn.connected() and 'UP' or 'DOWN')))
def _on_string_message(self, tcp_conn, message, receive_time):
logging.info('receive scan job %s' %message)
reply_msg = '%s received scan request %s' %(self._hostname, message)
self._codec.send(tcp_conn, reply_msg)
tcp_conn.shutdown()
root_xml = ElementTree.fromstring(message)
assert root_xml.tag == 'job'
if root_xml.attrib.get('name', None) == 'scan':
self._handle_scan(root_xml.text.strip())
def _handle_scan(self, scan_list):
tid = current_thread().ident
if not tid:
randome.random(time.time())
tid = random.randint(1, 1000000000)
root_list_file = 'scan_list_%d_%s' %(tid, strftime('%Y%m%d-%H%M%S'))
root_list_file = join(self._scan_list_dir, root_list_file)
with open(root_list_file, 'w') as scan_file:
scan_file.write(scan_list)
scan_cmd = ['smController', '-rlist', root_list_file] + self._scan_parameters
logging.info('executing %s' %' '.join(scan_cmd))
try:
output = subprocess.Popen(scan_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE).communicate()
result_msg = ''
if output[0]:
logging.info(output[0])
result_msg = output[0]
if output[1]:
logging.error(output[1])
result_msg += output[1]
except Exception as e:
exc_type, exc_value, exc_traceback = sys.exc_info()
result_msg = ''.join(format_exception(exc_type, exc_value, exc_traceback))
logging.error('Scanner._handle_scan has exception %s' %result_msg)
#raise
#self._codec.send(tcp_conn, result_msg)
def _construct_scan_parameters(self):
self._scan_parameters = []
stormax_section = 'STORMAX'
if stormax_section in self._cfg_parser.sections():
thrd_num = self._cfg_parser.get(stormax_section, 'thread_num')
if thrd_num:
self._scan_parameters.extend(['-thrds', thrd_num])
else:
self._scan_parameters.extend(['-thrds', str(cpu_count()*4)])
scan_cfg = self._cfg_parser.get(stormax_section, 'scan_cfg')
if scan_cfg:
self._scan_parameters.extend(['-scancfg', scan_cfg])
prescan = self._cfg_parser.get(stormax_section, 'prescan')
if prescan:
self._scan_parameters.extend(['-prescan', prescan])
postscan = self._cfg_parser.get(stormax_section, 'postscan')
if postscan:
self._scan_parameters.extend(['-postscan', postscan])
dtldir = self._cfg_parser.get(stormax_section, 'dtldir')
if dtldir:
self._scan_parameters.extend(['-dtldir', dtldir])
logdir = self._cfg_parser.get(stormax_section, 'logdir')
if logdir:
self._scan_parameters.extend(['-logdir', logdir])
errdir = self._cfg_parser.get(stormax_section, 'errdir')
if errdir:
self._scan_parameters.extend(['-errdir', errdir])
tmpdir = self._cfg_parser.get(stormax_section, 'tmpdir')
if tmpdir:
self._scan_parameters.extend(['-tmpdir', tmpdir])
