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)
