Updated tests for AIP blocklist generation.

aip.py

@@ -80,7 +80,7 @@ def find_new_data_files(self, raw_data_dir_path, processed_files_filepath) -> tu
         files_dates = {}
 
         try:
-            with open(processed_files_filepath, 'r') as records_file:
+            with open(processed_files_filepath, "r", encoding="utf-8") as records_file:
                 processed_files = records_file.read().split('\n')
         except IOError as e:
             logger.error(f"Unable to open {processed_files_filepath} file: {e}")
@@ -93,7 +93,7 @@ def find_new_data_files(self, raw_data_dir_path, processed_files_filepath) -> tu
 
         try:
             for new_file in new_data_files:
-                with open(processed_files_filepath, 'a') as records_file:
+                with open(processed_files_filepath, "a", encoding="utf-8") as records_file:
                     records_file.write(new_file + '\n')
         except (IOError, ValueError) as e:
             logger.error(f"Unable to update {processed_files_filepath} file: {e}")
@@ -120,7 +120,8 @@ def open_sort_new_file(self, raw_data_dir_path, new_files) -> tuple:
         new_ip_flows = []
         new_ips = []
         for file in new_files:
-            with open(f"{raw_data_dir_path}/{file}", 'r') as csv_file:
+            file_path = f"{raw_data_dir_path}/{file}"
+            with open(file_path, "r", encoding="utf-8") as csv_file:
                 for line in csv.reader(csv_file):
                     if line[0] != 'SrcAddr':
                         new_ip_flows.append(Flow.from_line(line))
@@ -335,7 +336,7 @@ def create_final_blocklist(self, blocklist_config) -> None:
         row_fieldnames = ('# Number', 'IP address', 'Rating')
 
         try:
-            with open(blocklist_filepath, 'wt', newline ='') as blocklist_file:
+            with open(blocklist_filepath, "wt", newline ='', encoding="utf-8") as blocklist_file:
                 header_writer = csv.DictWriter(blocklist_file, fieldnames=header_fieldnames)
                 header_writer.writeheader()
                 csv_writer = csv.DictWriter(blocklist_file, fieldnames=row_fieldnames)
@@ -427,7 +428,7 @@ def get_chosen_functions(self) -> list:
         """
         chosen_functions = []
         try:
-            with open(self.functions_filepath, 'r') as csv_file:
+            with open(self.functions_filepath, "r", encoding="utf-8") as csv_file:
                 for line in csv.reader(csv_file):
                     if line:
                         chosen_functions.extend(line)

data_sorter/main.py

@@ -81,8 +81,10 @@
 
 list_of_dictionaries = [dataset.get(key) for key in dataset.keys()]
 
-labels = {"SrcAddr", "total_events", "total_duration", "average_duration", "total_bytes", "average_bytes", "total_packets",
-          "average_packets", "first_event_time", "last_event_time"}
+labels = {"SrcAddr", "total_events", "total_duration", "average_duration",
+          "total_bytes", "average_bytes", "total_packets", "average_packets",
+          "last_event_time", "first_event_time"
+}
 
 try:
     with open(output, "w", encoding=Defaults.UTF_8.value) as f:

main/select_modules.py

@@ -11,9 +11,10 @@
 import os
 
 from methodology import Methodology
+from slips_aip_constants.defaults import EnvVars, Functions
 
 
-FILE_FOR_FUNCTIONS = os.environ['output_folder'] + '/selected_modules.csv'
+FILE_FOR_FUNCTIONS = f"{os.environ[EnvVars.OUTPUT_FOLDER.value]}/selected_modules.csv"
 
 manual = 0
 automatic = 1
@@ -57,9 +58,9 @@
 else:
 
     list_of_functions_that_were_choosen = [
-        'prioritize_consistent_normalized_ips',
-        'prioritize_new_normalized_ips',
-        'prioritize_only_normalized_today_ips'
+        Functions.PCN.value,
+        Functions.PNN.value,
+        Functions.POTN.value
     ]
 
 with open(FILE_FOR_FUNCTIONS, "w", encoding="utf-8") as file:

setup.py

@@ -20,7 +20,7 @@
     packages=find_packages(exclude=("tests",)),
     version=version,
     author='Stratosphere IPS',
-    author_email='stratosphereips@agents.fel.cvut.cz',
+    author_email='stratosphere@aic.fel.cvut.cz',
     url='https://github.com/stratosphereips/AIP-Blocklist-Algorithm',
     download_url=f"https://github.com/stratosphereips/AIP-Blocklist-Algorithm/tarball/{version}",
     description='The Attacker IP Prioritizer (AIP) algorithm is a IPv4 address blocklist generator.',

tests/mock_data/output_data/placeholder_resulting.csv

@@ -0,0 +1,2 @@
+IP, Score
+8.8.8.8, 10

tests/test_aip.py

@@ -7,14 +7,19 @@
 """
 #! /usr/local/bin/python3
 
+import glob
 import os
 import pytest
 from datetime import datetime
+from logging import getLogger
 
 from aip import AIP
 from slips_aip_constants.defaults import Blocklists, Defaults, EnvVars, Functions
 
 
+logger = getLogger(__name__)
+
+
 MOCK_DATA_DIR = "/tests/mock_data/"
 
 os.environ[EnvVars.OUTPUT_FOLDER.value] = MOCK_DATA_DIR
@@ -31,9 +36,9 @@ def get_number_of_ips(filepath):
     try:
         with open(filepath, "r", encoding="utf-8") as abs_data:
             number_of_lines = len(abs_data.readlines())
-        print(f"Number of lines in absolute data {number_of_lines}\n")
+        logger.info(f"Number of lines in absolute data {number_of_lines}\n")
     except IOError as e:
-        print(f"Unknown number of lines for {filepath}: {e}\n")
+        logger.exception(f"Unknown number of lines for {filepath}: {e}\n")
 
     return number_of_lines