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)
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])