Skip to content

Commit 34ca085

Browse files
A-atmosA-atmos
committed
Refactored modules according to the changes made
1 parent 6f059e0 commit 34ca085

File tree

16 files changed

+48
-37
lines changed

16 files changed

+48
-37
lines changed

‎modules/CESNET/CESNET.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -315,7 +315,7 @@ def run(self):
315315

316316
# in case of an interface or a file, push every time we get an alert
317317
if (
318-
utils.is_msg_intended_for(message, 'export_evidence')
318+
__database__.is_msg_intended_for(message, 'export_evidence')
319319
and self.send_to_warden
320320
):
321321
evidence = json.loads(message['data'])

‎modules/RiskIQ/RiskIQ.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ def run(self):
106106
self.shutdown_gracefully()
107107
return True
108108

109-
if utils.is_msg_intended_for(message, 'new_ip'):
109+
if __database__.is_msg_intended_for(message, 'new_ip'):
110110
ip = message['data']
111111
if utils.is_ignored_ip(ip):
112112
continue

‎modules/arp/arp.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -397,11 +397,17 @@ def run(self):
397397
self.arp_ts = time.time()
398398

399399
message = __database__.get_message(self.c1)
400+
# if message and 'stop_process' in message['data']:
401+
# print(f"ARP message for {message['data']}")
400402
if message and message['data'] == 'stop_process':
403+
with open('/home/ac/Desktop/workspace/message.txt', 'w') as file:
404+
file.write(message['data'])
405+
file.close()
406+
401407
self.shutdown_gracefully()
402408
return True
403409

404-
if utils.is_msg_intended_for(message, 'new_arp'):
410+
if __database__.is_msg_intended_for(message, 'new_arp'):
405411
flow = json.loads(message['data'])
406412
ts = flow['ts']
407413
profileid = flow['profileid']
@@ -454,7 +460,7 @@ def run(self):
454460
self.shutdown_gracefully()
455461
return True
456462

457-
if utils.is_msg_intended_for(message, 'tw_closed'):
463+
if __database__.is_msg_intended_for(message, 'tw_closed'):
458464
profileid_tw = message['data']
459465
# when a tw is closed, this means that it's too old so we don't check for arp scan in this time
460466
# range anymore

‎modules/blocking/blocking.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -342,7 +342,7 @@ def run(self):
342342
self.shutdown_gracefully()
343343
return True
344344
# There's an IP that needs to be blocked
345-
if utils.is_msg_intended_for(message, 'new_blocking'):
345+
if __database__.is_msg_intended_for(message, 'new_blocking'):
346346
# message['data'] in the new_blocking channel is a dictionary that contains
347347
# the ip and the blocking options
348348
# Example of the data dictionary to block or unblock an ip:

‎modules/exporting_alerts/exporting_alerts.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -324,7 +324,7 @@ def run(self):
324324
self.shutdown_gracefully()
325325
return True
326326

327-
if utils.is_msg_intended_for(msg, 'export_evidence'):
327+
if __database__.is_msg_intended_for(msg, 'export_evidence'):
328328
evidence = json.loads(msg['data'])
329329
description = evidence['description']
330330
if 'slack' in self.export_to and hasattr(self, 'BOT_TOKEN'):

‎modules/flowalerts/flowalerts.py‎

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1725,7 +1725,7 @@ def run(self):
17251725
if message and message['data'] == 'stop_process':
17261726
self.shutdown_gracefully()
17271727
return True
1728-
if utils.is_msg_intended_for(message, 'new_flow'):
1728+
if __database__.is_msg_intended_for(message, 'new_flow'):
17291729
new_flow = json.loads(message['data'])
17301730
profileid = new_flow['profileid']
17311731
twid = new_flow['twid']
@@ -1877,7 +1877,7 @@ def run(self):
18771877
if message and message['data'] == 'stop_process':
18781878
self.shutdown_gracefully()
18791879
return True
1880-
if utils.is_msg_intended_for(message, 'new_ssh'):
1880+
if __database__.is_msg_intended_for(message, 'new_ssh'):
18811881
data = message['data']
18821882
data = json.loads(data)
18831883
profileid = data['profileid']
@@ -1914,7 +1914,7 @@ def run(self):
19141914
if message and message['data'] == 'stop_process':
19151915
self.shutdown_gracefully()
19161916
return True
1917-
if utils.is_msg_intended_for(message, 'new_notice'):
1917+
if __database__.is_msg_intended_for(message, 'new_notice'):
19181918
data = message['data']
19191919
# Convert from json to dict
19201920
data = json.loads(data)
@@ -1970,7 +1970,7 @@ def run(self):
19701970
if message and message['data'] == 'stop_process':
19711971
self.shutdown_gracefully()
19721972
return True
1973-
if utils.is_msg_intended_for(message, 'new_ssl'):
1973+
if __database__.is_msg_intended_for(message, 'new_ssl'):
19741974
# Check for self signed certificates in new_ssl channel (ssl.log)
19751975
data = message['data']
19761976
# Convert from json to dict
@@ -2029,11 +2029,17 @@ def run(self):
20292029

20302030

20312031
message = __database__.get_message(self.c5)
2032+
# if message and 'stop_process' in message['data']:
2033+
20322034
if message and message['data'] == 'stop_process':
2035+
with open('/home/ac/Desktop/workspace/message.txt', 'w') as file:
2036+
file.write(message['data'])
2037+
file.close()
2038+
20332039
self.shutdown_gracefully()
20342040
return True
20352041

2036-
if utils.is_msg_intended_for(message, 'tw_closed'):
2042+
if __database__.is_msg_intended_for(message, 'tw_closed'):
20372043
profileid_tw = message['data'].split('_')
20382044
profileid, twid = f'{profileid_tw[0]}_{profileid_tw[1]}', profileid_tw[-1]
20392045
self.detect_data_upload_in_twid(profileid, twid)
@@ -2043,7 +2049,7 @@ def run(self):
20432049
if message and message['data'] == 'stop_process':
20442050
self.shutdown_gracefully()
20452051
return True
2046-
if utils.is_msg_intended_for(message, 'new_dns_flow'):
2052+
if __database__.is_msg_intended_for(message, 'new_dns_flow'):
20472053
data = json.loads(message['data'])
20482054
profileid = data['profileid']
20492055
twid = data['twid']
@@ -2086,7 +2092,7 @@ def run(self):
20862092
if message and message['data'] == 'stop_process':
20872093
self.shutdown_gracefully()
20882094
return True
2089-
if utils.is_msg_intended_for(message, 'new_downloaded_file'):
2095+
if __database__.is_msg_intended_for(message, 'new_downloaded_file'):
20902096
ssl_info = json.loads(message['data'])
20912097
self.check_malicious_ssl(ssl_info)
20922098

@@ -2095,7 +2101,7 @@ def run(self):
20952101
if message and message['data'] == 'stop_process':
20962102
self.shutdown_gracefully()
20972103
return True
2098-
if utils.is_msg_intended_for(message, 'new_smtp'):
2104+
if __database__.is_msg_intended_for(message, 'new_smtp'):
20992105
data = json.loads(message['data'])
21002106
profileid = data['profileid']
21012107
twid = data['twid']
@@ -2120,7 +2126,7 @@ def run(self):
21202126
if message and message['data'] == 'stop_process':
21212127
self.shutdown_gracefully()
21222128
return True
2123-
if utils.is_msg_intended_for(message, 'new_software'):
2129+
if __database__.is_msg_intended_for(message, 'new_software'):
21242130
flow = json.loads(message['data'])
21252131
starttime = flow.get('starttime', '')
21262132
saddr = flow.get('saddr', '')
@@ -2156,7 +2162,7 @@ def run(self):
21562162
self.shutdown_gracefully()
21572163
return True
21582164

2159-
if utils.is_msg_intended_for(message, 'new_weird'):
2165+
if __database__.is_msg_intended_for(message, 'new_weird'):
21602166
msg = json.loads(message['data'])
21612167
self.check_weird_http_method(msg)
21622168

@@ -2165,7 +2171,7 @@ def run(self):
21652171
self.shutdown_gracefully()
21662172
return True
21672173

2168-
if utils.is_msg_intended_for(message, 'new_tunnel'):
2174+
if __database__.is_msg_intended_for(message, 'new_tunnel'):
21692175
msg = json.loads(message['data'])
21702176
self.check_GRE_tunnel(msg)
21712177
except KeyboardInterrupt:

‎modules/flowmldetection/flowmldetection.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -393,7 +393,7 @@ def run(self):
393393
self.shutdown_gracefully()
394394
return True
395395

396-
if utils.is_msg_intended_for(message, 'new_flow'):
396+
if __database__.is_msg_intended_for(message, 'new_flow'):
397397
data = message['data']
398398
# Convert from json to dict
399399
data = json.loads(data)
@@ -441,7 +441,6 @@ def run(self):
441441
# Predict
442442
pred = self.detect()
443443
label = self.flow_dict['label']
444-
445444
# Report
446445
if (
447446
label

‎modules/http_analyzer/http_analyzer.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -471,7 +471,7 @@ def run(self):
471471
self.shutdown_gracefully()
472472
return True
473473

474-
if utils.is_msg_intended_for(message, 'new_http'):
474+
if __database__.is_msg_intended_for(message, 'new_http'):
475475
message = json.loads(message['data'])
476476
profileid = message['profileid']
477477
twid = message['twid']

‎modules/ip_info/ip_info.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -486,7 +486,7 @@ def run(self):
486486
if message and message['data'] == 'stop_process':
487487
self.shutdown_gracefully()
488488
return True
489-
if utils.is_msg_intended_for(message, 'new_MAC'):
489+
if __database__.is_msg_intended_for(message, 'new_MAC'):
490490
data = json.loads(message['data'])
491491
mac_addr = data['MAC']
492492
host_name = data.get('host_name', False)
@@ -509,7 +509,7 @@ def run(self):
509509
self.shutdown_gracefully()
510510
return True
511511

512-
if utils.is_msg_intended_for(message, 'new_dns_flow'):
512+
if __database__.is_msg_intended_for(message, 'new_dns_flow'):
513513
data = message['data']
514514
data = json.loads(data)
515515
# profileid = data['profileid']
@@ -528,7 +528,7 @@ def run(self):
528528
self.shutdown_gracefully()
529529
return True
530530

531-
if utils.is_msg_intended_for(message, 'new_ip'):
531+
if __database__.is_msg_intended_for(message, 'new_ip'):
532532
# Get the IP from the message
533533
ip = message['data']
534534
try:

‎modules/network_discovery/network_discovery.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -717,7 +717,7 @@ def run(self):
717717
self.shutdown_gracefully()
718718
return True
719719

720-
if utils.is_msg_intended_for(message, 'tw_modified'):
720+
if __database__.is_msg_intended_for(message, 'tw_modified'):
721721
# Get the profileid and twid
722722
profileid = message['data'].split(':')[0]
723723
twid = message['data'].split(':')[1]
@@ -751,7 +751,7 @@ def run(self):
751751
self.shutdown_gracefully()
752752
return True
753753

754-
if utils.is_msg_intended_for(message, 'new_notice'):
754+
if __database__.is_msg_intended_for(message, 'new_notice'):
755755
data = message['data']
756756
if type(data) != str:
757757
continue
@@ -776,7 +776,7 @@ def run(self):
776776
self.shutdown_gracefully()
777777
return True
778778

779-
if utils.is_msg_intended_for(message, 'new_dhcp'):
779+
if __database__.is_msg_intended_for(message, 'new_dhcp'):
780780
flow = json.loads(message['data'])
781781
self.check_dhcp_scan(flow)
782782

0 commit comments

Comments
 (0)