Ejemplo n.º 1
0
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)
Ejemplo n.º 2
0
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])