def test_011_fire_event(self): event_name = 'newCoolEvent' eh = EventHandler(event_name) def callback1(*args, **kwargs): self.assertEqual(args[0], 1) self.assertEqual(args[1], 2) self.assertEqual(args[2], 3) self.assertEqual(kwargs['extra'], 0) self.assertTrue(eh.link(callback1, event_name)) self.assertTrue(eh.fire(event_name, 1, 2, 3, extra=0)) def will_fail_callback(number1, number2, number3, extra=0): return number1 / extra self.assertTrue(eh.link(will_fail_callback, event_name)) with self.assertRaises(ZeroDivisionError) as context: eh.fire(event_name, 1, 2, 3, extra=0) # Set callback fail toleration eh.verbose = True eh.tolerate_exceptions = True output = io.StringIO() eh.stream_output = output self.assertFalse(eh.fire(event_name, 1, 2, 3, extra=0)) value = output.getvalue() self.assertTrue('WARNING' in value) self.assertTrue(will_fail_callback.__name__ in value)
class ChatRoom: """Simulates a chatroom environment with event handler implementation. This is just a documented sample without pretensions. It is not a real class implementation. """ def __init__(self): """Initialize the chat room.""" self.__messages = [] # Stores users messages self.__users = { 'bot': [] } # Stores a dictionary with registered usernames # Define the event handler and make it public outside the class to let externals subscriptions to events. self.event_handler = EventHandler( 'onNewuser', 'onMessage') # Note that events names are cased sensitive. # You can set any number of unique events and asing any number of unique callbacks to fire per event. # Is not necessary define events names during initialization, also you can register the event names during # run time using register_event method. # Lets link some internal class methods to those events as callbacks. # Limits are available resources. self.event_handler.link(self.__on_newuser_join, 'onNewuser') self.event_handler.link(self.__on_message, 'onMessage') # Now lets define this two methods to dispatch the events # Note this methods are not accesible outside class instance # This calbback will be called when onNewUser event happens def __on_newuser_join(self, user): """Shout in the output telling new user has joined chat room, when onNewuser event happens.""" print( f'\n** ChatRoom info ** user {user} has joined the chat ** {len(self.user_list())} user/s **\n' ) # This callback will be called when onMessage event happens def __on_message(self, user, msg): """Print the user message in the output, when onMessage event happens.""" print(f'{user} says:\t {msg}') # Now let's define the public methods of the chatroom to be used outside the class def user_list(self): """Return a list of not bot users.""" return [user for user in self.__users.keys() if user != 'bot'] def say(self, user, msg=None): """Let user (and bots) send a message to the chat room.""" if not user in self.__users: # if user is not registered fire onNewuser event and recibe it inside the class. self.__users[user] = [] self.event_handler.fire('onNewuser', user) if not msg: return if msg != '': # Enqueue the message and fire onMessage event to be received internally by __on_message method. self.__messages.append((user, msg)) self.event_handler.fire('onMessage', user, msg)
class Session(object): def __init__(self): self.predictor = Predictor() self.notificator = EventHandler('add_notify') def start(self): self.run_flag = True while self.run_flag: self.session_handler() def add_notificator(self, func): self.notificator.link(func, 'add_notify') def remove_notificator(self, func): pass def stop(self): self.run_flag = False def check_session(self, client, ip, id): if client['pkt_count'] > 5 and client['susp_pkt'] / client[ 'pkt_count'] > 0.1: db_susp_session.insert_one({"ses_id": id, "ip": ip}) self.notificator.fire( 'add_notify', 'Обнаружено туннелирование с адреса ' + str(ip).replace('_', '.')) return 1 return 0 def stop_session(self): self.end_time = datetime.datetime.now() def session_handler(self): start_time = datetime.datetime.now() self.end_time = start_time + datetime.timedelta(minutes=1) self.session_doc = { 'start_time': start_time, 'end_time': '0', 'clients': {}, 'dns_names': {} } id = db_sessions_col.insert_one(self.session_doc).inserted_id last = db_pkt_col_packets.find().sort( '$natural', pymongo.DESCENDING).limit(-1).next() last_id = last['_id'] while datetime.datetime.now() < self.end_time: cursor = db_pkt_col_packets.find({'_id': { '$gt': last_id }}) # ,cursor_type=pymongo.CursorType.TAILABLE_AWAIT) # while cursor.alive: for doc in cursor: pred = self.predictor.predict_pkt(doc) doc["label"] = pred doc['ses_id'] = id doc.pop('_id', None) db_archive_pkt_col.insert_one(doc) last_id = doc['_id'] ip = doc['src'] if doc['qr'] == '0' else doc['dst'] ip = str(ip).replace('.', '_') self.session_doc['clients'].update({ ip: { 'pkt_count': 1, 'susp_pkt': 0, 'susp_event': 0 } if (self.session_doc['clients'].get(ip) is None) else { 'pkt_count': self.session_doc['clients'][ip]['pkt_count'] + 1, 'susp_pkt': self.session_doc['clients'][ip]['susp_pkt'] + pred, 'susp_event': self.check_session(self.session_doc['clients'][ip], ip, id) if (self.session_doc['clients'][ip].get('susp_event') == 0) else 1 } }) qry_name = '.'.join(doc['qry_name'].split('.')[-2:]) qry_name = qry_name.replace('.', '_') self.session_doc['dns_names'].update({ qry_name: { 'count_queries': 1, 'subdomains': [doc['qry_name']] } if (self.session_doc['dns_names'].get(qry_name) is None) else { 'count_queries': self.session_doc['dns_names'][qry_name] ['count_queries'] + 1, 'subdomains': self.session_doc['dns_names'][qry_name]['subdomains'] if (doc['qry_name'] in self.session_doc['dns_names'] [qry_name]['subdomains']) else self.session_doc['dns_names'][qry_name]['subdomains'] + [doc['qry_name']] } }) db_sessions_col.update_one({'_id': id}, {'$set': self.session_doc}, upsert=False) time.sleep(1) self.session_doc['end_time'] = self.end_time db_sessions_col.update_one({'_id': id}, {'$set': self.session_doc}, upsert=False)
class ModuleHandler(): def __init__(self, bot): self.event_handler = EventHandler() self.bot = bot self.load_modules(True) def load_modules(self, first_time = False): if any(result is False for result in self.fire_event('modulehandler:before_load_modules', (self, self.bot, self.event_handler, first_time))): return False logging.info('Loading module resources') module_files = glob.glob(os.path.dirname(__file__) + '/modules/resources/*.py') module_names = ['modules.resources.' + os.path.basename(f)[:-3] for f in module_files] if not self._load_module_list(module_names, first_time): return False logging.info('Loading modules') module_files = glob.glob(os.path.dirname(__file__) + '/modules/*.py') module_names = ['modules.' + os.path.basename(f)[:-3] for f in module_files] imported_modules = self._load_module_list(module_names, first_time) if not imported_modules: return False if any(result is False for result in self.fire_event('modulehandler:before_init_modules', (self, self.bot, self.event_handler, first_time))): return False # we've imported with no problems - break old hooks, and try to add new self.event_handler.clear_module_hooks() self.event_handler.importing_modules = True for module in imported_modules: try: module.init() except BaseException as e: if first_time: raise # if we fail out here, this module will NOT have hooked its events - others may have though, so we don't return so all unbroken modules can init error = 'module "%s" unable to init: %s: %s' % (str(module), type(e).__name__, e) logging.exception(error) print(error) self.event_handler.importing_modules = False self.fire_event('modulehandler:after_load_modules', (self, self.bot, self.event_handler, first_time)) return True def _load_module_list(self, list, first_time): loaded_modules = [] for module in list: if module.endswith('__init__'): continue if module in sys.modules: module = sys.modules[module] importfunc = reload_func else: importfunc = importlib.import_module try: module = importfunc(module) module.bot = self.bot module.event_handler = self.event_handler loaded_modules.append(module) except BaseException as e: if first_time: raise # if we fail out here, old event hooks remain in place error = 'module "%s" unable to import: %s: %s' % (str(module), type(e).__name__, e) logging.exception(error) print(error) return [] return loaded_modules def fire_event(self, key, parameters): return self.event_handler.fire(key, parameters) def hook_event(self, key, function, priority = 500): return self.event_handler.hook(key, function, priority) def get_event_handlers(self, key): return self.event_handler.get_handlers(key)