Refactored modules according to the changes made

Author: A-atmos

modules/CESNET/CESNET.py

@@ -315,7 +315,7 @@ def run(self):
 
                 # in case of an interface or a file, push every time we get an alert
                 if (
-                    utils.is_msg_intended_for(message, 'export_evidence')
+                    __database__.is_msg_intended_for(message, 'export_evidence')
                     and self.send_to_warden
                 ):
                     evidence = json.loads(message['data'])

modules/RiskIQ/RiskIQ.py

@@ -106,7 +106,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'new_ip'):
+                if __database__.is_msg_intended_for(message, 'new_ip'):
                     ip = message['data']
                     if utils.is_ignored_ip(ip):
                         continue

modules/arp/arp.py

@@ -397,11 +397,17 @@ def run(self):
                     self.arp_ts = time.time()
 
                 message = __database__.get_message(self.c1)
+                # if message and 'stop_process' in message['data']:
+                #     print(f"ARP message for {message['data']}")
                 if message and message['data'] == 'stop_process':
+                    with open('/home/ac/Desktop/workspace/message.txt', 'w') as file:
+                        file.write(message['data'])
+                        file.close()
+                        
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'new_arp'):
+                if __database__.is_msg_intended_for(message, 'new_arp'):
                     flow = json.loads(message['data'])
                     ts = flow['ts']
                     profileid = flow['profileid']
@@ -454,7 +460,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'tw_closed'):
+                if __database__.is_msg_intended_for(message, 'tw_closed'):
                     profileid_tw = message['data']
                     # when a tw is closed, this means that it's too old so we don't check for arp scan in this time
                     # range anymore

modules/blocking/blocking.py

@@ -342,7 +342,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
                 # There's an IP that needs to be blocked
-                if utils.is_msg_intended_for(message, 'new_blocking'):
+                if __database__.is_msg_intended_for(message, 'new_blocking'):
                     # message['data'] in the new_blocking channel is a dictionary that contains
                     # the ip and the blocking options
                     # Example of the data dictionary to block or unblock an ip:

modules/exporting_alerts/exporting_alerts.py

@@ -324,7 +324,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(msg, 'export_evidence'):
+                if __database__.is_msg_intended_for(msg, 'export_evidence'):
                     evidence = json.loads(msg['data'])
                     description = evidence['description']
                     if 'slack' in self.export_to and hasattr(self, 'BOT_TOKEN'):

modules/flowalerts/flowalerts.py

@@ -1725,7 +1725,7 @@ def run(self):
                 if message and message['data'] == 'stop_process':
                     self.shutdown_gracefully()
                     return True
-                if utils.is_msg_intended_for(message, 'new_flow'):
+                if __database__.is_msg_intended_for(message, 'new_flow'):
                     new_flow = json.loads(message['data'])
                     profileid = new_flow['profileid']
                     twid = new_flow['twid']
@@ -1877,7 +1877,7 @@ def run(self):
                 if message and message['data'] == 'stop_process':
                     self.shutdown_gracefully()
                     return True
-                if utils.is_msg_intended_for(message, 'new_ssh'):
+                if __database__.is_msg_intended_for(message, 'new_ssh'):
                     data = message['data']
                     data = json.loads(data)
                     profileid = data['profileid']
@@ -1914,7 +1914,7 @@ def run(self):
                 if message and message['data'] == 'stop_process':
                     self.shutdown_gracefully()
                     return True
-                if utils.is_msg_intended_for(message, 'new_notice'):
+                if __database__.is_msg_intended_for(message, 'new_notice'):
                     data = message['data']
                     # Convert from json to dict
                     data = json.loads(data)
@@ -1970,7 +1970,7 @@ def run(self):
                 if message and message['data'] == 'stop_process':
                     self.shutdown_gracefully()
                     return True
-                if utils.is_msg_intended_for(message, 'new_ssl'):
+                if __database__.is_msg_intended_for(message, 'new_ssl'):
                     # Check for self signed certificates in new_ssl channel (ssl.log)
                     data = message['data']
                     # Convert from json to dict
@@ -2029,11 +2029,17 @@ def run(self):
 
 
                 message = __database__.get_message(self.c5)
+                # if message and 'stop_process' in message['data']:
+
                 if message and message['data'] == 'stop_process':
+                    with open('/home/ac/Desktop/workspace/message.txt', 'w') as file:
+                        file.write(message['data'])
+                        file.close()
+                        
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'tw_closed'):
+                if __database__.is_msg_intended_for(message, 'tw_closed'):
                     profileid_tw = message['data'].split('_')
                     profileid, twid = f'{profileid_tw[0]}_{profileid_tw[1]}', profileid_tw[-1]
                     self.detect_data_upload_in_twid(profileid, twid)
@@ -2043,7 +2049,7 @@ def run(self):
                 if message and message['data'] == 'stop_process':
                     self.shutdown_gracefully()
                     return True
-                if utils.is_msg_intended_for(message, 'new_dns_flow'):
+                if __database__.is_msg_intended_for(message, 'new_dns_flow'):
                     data = json.loads(message['data'])
                     profileid = data['profileid']
                     twid = data['twid']
@@ -2086,7 +2092,7 @@ def run(self):
                 if message and message['data'] == 'stop_process':
                     self.shutdown_gracefully()
                     return True
-                if utils.is_msg_intended_for(message, 'new_downloaded_file'):
+                if __database__.is_msg_intended_for(message, 'new_downloaded_file'):
                     ssl_info = json.loads(message['data'])
                     self.check_malicious_ssl(ssl_info)
 
@@ -2095,7 +2101,7 @@ def run(self):
                 if message and message['data'] == 'stop_process':
                     self.shutdown_gracefully()
                     return True
-                if utils.is_msg_intended_for(message, 'new_smtp'):
+                if __database__.is_msg_intended_for(message, 'new_smtp'):
                     data = json.loads(message['data'])
                     profileid = data['profileid']
                     twid = data['twid']
@@ -2120,7 +2126,7 @@ def run(self):
                 if message and message['data'] == 'stop_process':
                     self.shutdown_gracefully()
                     return True
-                if utils.is_msg_intended_for(message, 'new_software'):
+                if __database__.is_msg_intended_for(message, 'new_software'):
                     flow = json.loads(message['data'])
                     starttime = flow.get('starttime', '')
                     saddr = flow.get('saddr', '')
@@ -2156,7 +2162,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'new_weird'):
+                if __database__.is_msg_intended_for(message, 'new_weird'):
                     msg = json.loads(message['data'])
                     self.check_weird_http_method(msg)
 
@@ -2165,7 +2171,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'new_tunnel'):
+                if __database__.is_msg_intended_for(message, 'new_tunnel'):
                     msg = json.loads(message['data'])
                     self.check_GRE_tunnel(msg)
             except KeyboardInterrupt:

modules/flowmldetection/flowmldetection.py

@@ -393,7 +393,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'new_flow'):
+                if __database__.is_msg_intended_for(message, 'new_flow'):
                     data = message['data']
                     # Convert from json to dict
                     data = json.loads(data)
@@ -441,7 +441,6 @@ def run(self):
                             # Predict
                             pred = self.detect()
                             label = self.flow_dict['label']
-
                             # Report
                             if (
                                 label

modules/http_analyzer/http_analyzer.py

@@ -471,7 +471,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'new_http'):
+                if __database__.is_msg_intended_for(message, 'new_http'):
                     message = json.loads(message['data'])
                     profileid = message['profileid']
                     twid = message['twid']

modules/ip_info/ip_info.py

@@ -486,7 +486,7 @@ def run(self):
                 if message and message['data'] == 'stop_process':
                     self.shutdown_gracefully()
                     return True
-                if utils.is_msg_intended_for(message, 'new_MAC'):
+                if __database__.is_msg_intended_for(message, 'new_MAC'):
                     data = json.loads(message['data'])
                     mac_addr = data['MAC']
                     host_name = data.get('host_name', False)
@@ -509,7 +509,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'new_dns_flow'):
+                if __database__.is_msg_intended_for(message, 'new_dns_flow'):
                     data = message['data']
                     data = json.loads(data)
                     # profileid = data['profileid']
@@ -528,7 +528,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'new_ip'):
+                if __database__.is_msg_intended_for(message, 'new_ip'):
                     # Get the IP from the message
                     ip = message['data']
                     try:

modules/network_discovery/network_discovery.py

@@ -717,7 +717,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'tw_modified'):
+                if __database__.is_msg_intended_for(message, 'tw_modified'):
                     # Get the profileid and twid
                     profileid = message['data'].split(':')[0]
                     twid = message['data'].split(':')[1]
@@ -751,7 +751,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'new_notice'):
+                if __database__.is_msg_intended_for(message, 'new_notice'):
                     data = message['data']
                     if type(data) != str:
                         continue
@@ -776,7 +776,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'new_dhcp'):
+                if __database__.is_msg_intended_for(message, 'new_dhcp'):
                     flow = json.loads(message['data'])
                     self.check_dhcp_scan(flow)
 

modules/p2ptrust/p2ptrust.py

@@ -635,23 +635,23 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'report_to_peers'):
+                if __database__.is_msg_intended_for(message, 'report_to_peers'):
                     self.new_evidence_callback(message)
 
                 message = __database__.get_message(self.c2)
                 if message and message['data'] == 'stop_process':
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, self.p2p_data_request_channel):
+                if __database__.is_msg_intended_for(message, self.p2p_data_request_channel):
                     self.data_request_callback(message)
 
                 message = __database__.get_message(self.c3)
                 if message and message['data'] == 'stop_process':
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, self.gopy_channel):
+                if __database__.is_msg_intended_for(message, self.gopy_channel):
                     self.gopy_callback(message)
 
                 ret_code = self.pigeon.poll()

modules/rnn-cc-detection/rnn-cc-detection.py

@@ -140,7 +140,7 @@ def run(self, model_file='modules/rnn-cc-detection/rnn_model.h5'):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'new_letters'):
+                if __database__.is_msg_intended_for(message, 'new_letters'):
                     data = message['data']
                     data = json.loads(data)
                     pre_behavioral_model = data['new_symbol']

modules/threat_intelligence/threat_intelligence.py

@@ -946,7 +946,7 @@ def run(self):
                     self.should_shutdown = True
 
                 # The channel now can receive an IP address or a domain name
-                if utils.is_msg_intended_for(
+                if __database__.is_msg_intended_for(
                     message, 'give_threat_intelligence'
                 ):
                     # Data is sent in the channel as a json dict so we need to deserialize it first
@@ -1005,7 +1005,7 @@ def run(self):
                 if message and message['data'] == 'stop_process':
                     self.should_shutdown = True
 
-                if utils.is_msg_intended_for(message, 'new_downloaded_file'):
+                if __database__.is_msg_intended_for(message, 'new_downloaded_file'):
                     file_info = json.loads(message['data'])
                     self.is_malicious_hash(file_info)
 

modules/timeline/timeline.py

@@ -383,7 +383,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'new_flow'):
+                if __database__.is_msg_intended_for(message, 'new_flow'):
                     mdata = message['data']
                     # Convert from json to dict
                     mdata = json.loads(mdata)

modules/update_manager/update_manager.py

@@ -120,7 +120,7 @@ def run(self):
             try:
                 message = __database__.get_message(self.c1)
                 # Check that the message is for you. Probably unnecessary...
-                if message and message['data'] == 'stop_process':
+                if message and ('stop_process' in message['data']):
                     self.shutdown_gracefully()
                     return True
 

modules/virustotal/virustotal.py

@@ -535,7 +535,7 @@ def run(self):
                     self.shutdown_gracefully()
                     return True
 
-                if utils.is_msg_intended_for(message, 'new_flow'):
+                if __database__.is_msg_intended_for(message, 'new_flow'):
                     data = message['data']
                     data = json.loads(data)
                     # profileid = data['profileid']
@@ -576,7 +576,7 @@ def run(self):
                 if message and message['data'] == 'stop_process':
                     self.shutdown_gracefully()
                     return True
-                if utils.is_msg_intended_for(message, 'new_dns_flow'):
+                if __database__.is_msg_intended_for(message, 'new_dns_flow'):
                     data = message['data']
                     data = json.loads(data)
                     # profileid = data['profileid']
@@ -610,7 +610,7 @@ def run(self):
                 if message and message['data'] == 'stop_process':
                     self.shutdown_gracefully()
                     return True
-                if utils.is_msg_intended_for(message, 'new_url'):
+                if __database__.is_msg_intended_for(message, 'new_url'):
                     data = message['data']
                     data = json.loads(data)
                     # profileid = data['profileid']