Skip to content
GitLab
Commits (2)
Hide whitespace changes
Inline Side-by-side
src/l3_distributedattackdetector/service/l3_distributedattackdetector.py
View file @ 90be042f
......@@ -65,6 +65,18 @@ BATCH_SIZE = 10
class l3_distributedattackdetector:
def __init__(self):
"""
Initializes a Distributed Attack Detector.
This method initializes a Distributed Attack Detector by setting up instance variables, connecting to the Centralized Attack Detector, obtaining feature IDs, and starting the process traffic loop. It also sets up a signal handler for handling keyboard interrupts.
Args:
None.
Returns:
None.
"""
LOGGER.info("Creating Distributed Attack Detector")
self.feature_ids = []
......@@ -88,27 +100,67 @@ class l3_distributedattackdetector:
LOGGER.info("Features Ids.: {:s}".format(str(self.feature_ids)))
asyncio.run(self.process_traffic())
def read_kwnown_attack_ips(self):
"""
Reads a list of known attack IPs from a CSV file.
This method reads a list of known attack IPs from a CSV file named "known_attack_ips.csv". The method returns a list of strings containing the IP addresses.
Args:
None.
Returns:
List[str]: A list of strings containing the IP addresses of known attack sources.
"""
# Initialize an empty list to store the known attack IPs
known_attack_ips = []
# open known attack ips csv file
# Open the known attack IPs CSV file
with open("known_attack_ips.csv", "r") as f:
# Read the contents of the file and split it into a list of strings
known_attack_ips = f.read().split(",")
# Return the list of known attack IPs
return known_attack_ips
def handler(self):
"""
Handles a keyboard interrupt signal.
This method handles a keyboard interrupt signal by setting the `STOP` flag to `True` and logging a message indicating that the program is stopping gracefully.
Args:
None.
Returns:
None.
"""
# Set the STOP flag to True
if STOP:
exit()
STOP = True
# Log a message indicating that the program is stopping gracefully
LOGGER.info("Gracefully stopping...")
def follow(self, logfile, time_sleep):
"""
Generator function that yields new lines in a log file
Generator function that yields new lines in a log file.
This method reads a file object and yields new lines as they are added to the file. The method uses an infinite loop to continuously read the last line of the file. If the file hasn't been updated, the method sleeps for the specified `time_sleep` interval. If the last line of the file doesn't end with a newline character, the method appends the line to a `chunk` variable. When a newline character is encountered, the method yields the line, possibly appending the `chunk` variable to the line if it contains a partial line. The method returns an iterator that yields lines from the file.
Args:
file (TextIO): The file object to read from.
time_sleep (float): The time to sleep if the file hasn't been updated.
Yields:
str: The next line in the file.
Returns:
None.
"""
# seek the end of the file
......@@ -136,11 +188,27 @@ class l3_distributedattackdetector:
yield line
def load_file(self, dirname=TSTAT_DIR_NAME):
"""
Loads the latest Tstat log file.
This method loads the latest Tstat log file by searching for the most recent directory in the specified `dirname` directory. If a directory is found, the method returns the path to the `log_tcp_temp_complete` file in that directory. If no directory is found, the method logs a message and waits for 5 seconds before trying again.
Args:
dirname (str): The name of the directory to search for Tstat log files. Defaults to `TSTAT_DIR_NAME`.
Returns:
str: The path to the latest Tstat log file.
"""
while True:
# Get the path to the Tstat directory
here = os.path.dirname(os.path.abspath(__file__))
tstat_piped = os.path.join(here, dirname)
# Get a list of all directories in the Tstat directory
tstat_dirs = os.listdir(tstat_piped)
# If there are directories in the Tstat directory, find the most recent one and return the path to the log file
if len(tstat_dirs) > 0:
tstat_dirs.sort()
new_dir = tstat_dirs[-1]
......@@ -149,16 +217,25 @@ class l3_distributedattackdetector:
LOGGER.info("Following: {:s}".format(str(tstat_file)))
return tstat_file
# If there are no directories in the Tstat directory, log a message and wait for 5 seconds before trying again
else:
LOGGER.info("No Tstat directory found. Waiting...")
time.sleep(5)
def process_line(self, line):
"""
- Preprocessing before a message per line
- Avoids crash when nan are found by generating a 0s array
- Returns a list of values
Processes a single line of input data and returns a list of feature values.
Args:
line (str): A single line of input data containing feature values separated by spaces.
Returns:
List[float]: A list of feature values extracted from the input line.
Raises:
IndexError: If the input line does not contain enough feature values.
"""
line = line.split(" ")
try:
......@@ -172,15 +249,19 @@ class l3_distributedattackdetector:
return values
def get_service_ids(self, context_id_str):
with grpc.insecure_channel(CONTEXT_CHANNEL) as channel:
stub = ContextServiceStub(channel)
context_id = ContextId()
context_id.context_uuid.uuid = context_id_str
def get_services(self, context_id_str):
"""
Gets the services for a given context ID.
return stub.ListServiceIds(context_id)
This method gets the services for a given context ID by calling the ListServices method of the ContextService gRPC stub.
def get_services(self, context_id_str):
Args:
context_id_str (str): The context ID to get services for.
Returns:
ListServicesResponse: A response object containing the services.
"""
with grpc.insecure_channel(CONTEXT_CHANNEL) as channel:
stub = ContextServiceStub(channel)
context_id = ContextId()
......@@ -189,6 +270,18 @@ class l3_distributedattackdetector:
return stub.ListServices(context_id)
def get_service_id(self, context_id):
"""
Gets the service ID for a given context ID.
This method gets the service ID for a given context ID by calling the get_services method and searching for the first service in the list with the service type of SERVICETYPE_L3NM.
Args:
context_id (str): The context ID to get the service ID for.
Returns:
str: The service ID.
"""
service_list = self.get_services(context_id)
service_id = None
......@@ -202,6 +295,18 @@ class l3_distributedattackdetector:
return service_id
def get_endpoint_id(self, context_id):
"""
Gets the endpoint ID for a given context ID.
This method gets the endpoint ID for a given context ID by calling the get_services method and searching for the first service in the list with the service type of SERVICETYPE_L3NM.
Args:
context_id (str): The context ID to get the endpoint ID for.
Returns:
str: The endpoint ID.
"""
service_list = self.get_services(context_id)
endpoint_id = None
......@@ -213,9 +318,33 @@ class l3_distributedattackdetector:
return endpoint_id
def get_features_ids(self):
"""
Gets the feature IDs used by the Centralized Attack Detector model.
This method gets the feature IDs used by the Centralized Attack Detector model by calling the GetFeaturesIds method of the Centralized Attack Detector gRPC stub.
Args:
None.
Returns:
list: A list of feature IDs.
"""
return self.cad.GetFeaturesIds(Empty()).auto_features
def check_types(self):
"""
Checks the types of the features that will be sent to the Centralized Attack Detector.
This method checks the types of the Centralized Attack Detector features to ensure that they are of the correct type. If any of the types are incorrect, the method raises an AssertionError.
Args:
None.
Returns:
None.
"""
for feature in self.cad_features["features"]:
assert isinstance(feature, float)
......@@ -229,6 +358,28 @@ class l3_distributedattackdetector:
assert isinstance(self.cad_features["connection_metadata"]["time_end"], float)
def insert_connection(self):
"""
Inserts a new connection into the `connections_dict` instance variable.
This method inserts a new connection into the `connections_dict` instance variable. The method uses the `conn_id` instance variable to create a new dictionary entry for the connection. If the connection already exists in the dictionary, the method updates the `time_end` value of the existing entry. If the connection doesn't exist in the dictionary, the method creates a new entry with the following keys:
- "ip_o": The source IP address of the connection.
- "port_o": The source port of the connection.
- "ip_d": The destination IP address of the connection.
- "port_d": The destination port of the connection.
- "flow_id": The flow ID of the connection.
- "service_id": The service ID of the connection.
- "endpoint_id": The endpoint ID of the connection.
- "protocol": The protocol of the connection.
- "time_start": The start time of the connection.
- "time_end": The end time of the connection.
Args:
None.
Returns:
None.
"""
try:
self.connections_dict[self.conn_id]["time_end"] = time.time()
except KeyError:
......@@ -245,10 +396,46 @@ class l3_distributedattackdetector:
self.connections_dict[self.conn_id]["port_d"] = self.conn_id[3]
def check_if_connection_is_attack(self):
"""
Checks if a connection is an attack based on known attack IP addresses.
This method checks if a connection is an attack based on known attack IP addresses. The method uses the `conn_id` and `known_attack_ips` instance variables to determine if the source or destination IP address of the connection is in the list of known attack IP addresses. If either IP address is in the list, the method logs a message indicating that an attack has been detected.
Args:
None.
Returns:
None.
"""
if self.conn_id[0] in self.known_attack_ips or self.conn_id[2] in self.known_attack_ips:
LOGGER.info("Attack detected. Origin IP address: {0}, destination IP address: {1}".format(self.conn_id[0], self.conn_id[2]))
def create_cad_features(self):
"""
Creates a dictionary of features and connection metadata for the Centralized Attack Detector.
This method creates a dictionary of features and connection metadata for the Centralized Attack Detector. The method uses the `new_connections` and `connections_dict` instance variables to obtain the necessary data. The resulting dictionary contains the following keys:
- "features": A list of the first 10 features of the connection.
- "connection_metadata": A dictionary containing the following keys:
- "ip_o": The source IP address of the connection.
- "port_o": The source port of the connection.
- "ip_d": The destination IP address of the connection.
- "port_d": The destination port of the connection.
- "flow_id": The flow ID of the connection.
- "service_id": The service ID of the connection.
- "endpoint_id": The endpoint ID of the connection.
- "protocol": The protocol of the connection.
- "time_start": The start time of the connection.
- "time_end": The end time of the connection.
Args:
None.
Returns:
None.
"""
self.cad_features = {
"features": self.new_connections[self.conn_id][0:10],
"connection_metadata": {
......@@ -266,6 +453,18 @@ class l3_distributedattackdetector:
}
async def send_batch_async(self, metrics_list_pb):
"""
Sends a batch of traffic data to a Centralized Attack Detector.
This method sends a batch of traffic data to a Centralized Attack Detector for analysis. The method creates a `L3CentralizedattackdetectorBatchInput` object from the provided `metrics_list_pb`, and sends the batch using an executor to run the `AnalyzeBatchConnectionStatistics` method in a separate thread. The method returns `None`.
Args:
metrics_list_pb (List[L3CentralizedattackdetectorMetrics]): A list of traffic metrics to send to the Centralized Attack Detector.
Returns:
None.
"""
loop = asyncio.get_running_loop()
# Create metrics batch
......@@ -281,7 +480,19 @@ class l3_distributedattackdetector:
LOGGER.error(f"Error sending batch: {e}")
async def send_data(self, metrics_list_pb, send_data_times):
# Send data to CAD
"""
Sends traffic data to a Centralized Attack Detector.
This method sends traffic data to a Centralized Attack Detector for analysis. If the `SEND_DATA_IN_BATCHES` flag is set to `True`, the data is sent in batches of size `BATCH_SIZE`. Otherwise, the data is sent one metric at a time. The method returns the updated `metrics_list_pb` and `send_data_times` arrays.
Args:
metrics_list_pb (List[L3CentralizedattackdetectorMetrics]): A list of traffic metrics to send to the Centralized Attack Detector.
send_data_times (np.ndarray): An array of times it took to send each batch of data.
Returns:
Tuple[List[L3CentralizedattackdetectorMetrics], np.ndarray]: A tuple containing the updated `metrics_list_pb` and `send_data_times` arrays.
"""
if SEND_DATA_IN_BATCHES:
if len(metrics_list_pb) == BATCH_SIZE:
send_data_time_start = time.time()
......@@ -303,6 +514,15 @@ class l3_distributedattackdetector:
return metrics_list_pb, send_data_times
async def process_traffic(self):
""" Processes traffic data from a Tstat log file.
This method reads traffic data from a Tstat log file, processes each line of data, and sends the resulting metrics to a Centralized Attack Detector. It runs indefinitely until the `STOP` flag is set to `True`.
Args:
None.
Returns:
None.
"""
LOGGER.info("Loading Tstat log file...")
logfile = open(self.load_file(), "r")
......