Merge development into main

lib/aip/data/access.py

@@ -36,52 +36,94 @@
 import pandas as pd
 import numpy as np
 
-from aip.data.functions import scramble, read_zeek, getrawdata, removerawdata
-from joblib import Parallel, delayed
-from os import scandir, path
+from joblib import Parallel
+from joblib import delayed
+from os import scandir
+from os import path
 from pathlib import Path
+from aip.data.functions import scramble
+from aip.data.functions import read_zeek
+from aip.data.functions import get_raw_data
+from aip.data.functions import remove_raw_data
 
 project_dir = Path(__file__).resolve().parents[3]
 
 data_path = path.join(project_dir,'data')  # Deprecated, do not use
 data_dir = path.join(project_dir,'data')
 
+
 def _get_honeypot_ips(for_date=None):
     '''
-    Filter those honeypots active due date for_date, if there are operation dates in the honeypot file.
+    Return the public IPs of the honeypots. If for_date is given and the
+    honeypot file has operation dates, only return the honeypots active on
+    that date. The honeypots_public_ips.csv has the following format:
+        # List of IPs to look for to generate the attack files.
+        public_ip,operation_start_date,operation_end_date
     '''
-    logger = logging.getLogger(__name__)
     # Check if the file exists before attempting to read it
     honeypot_public_ips = path.join(project_dir, 'data', 'external', 'honeypots_public_ips.csv')
 
+    # If the file does not exist raise an exception
     if not path.exists(honeypot_public_ips):
-        logger.error(f"File 'honeypot_public_ips.csv' does not exist. Raising error.")
-        raise FileNotFoundError("Required file 'honeypots_public_ips.csv' does not exist.")
+        raise FileNotFoundError("_get_honeypot_ips() required file 'honeypots_public_ips.csv' does not exist.")
 
+    # Read CSV located in data/external/honeypots_public_ips.csv
     honeypots = pd.read_csv(path.join(project_dir, 'data', 'external', 'honeypots_public_ips.csv'), comment='#')
+
     if for_date is not None:
+        # Convert to datetime object
         for_date = pd.to_datetime(for_date)
+
+        # Parsing start date
         if 'operation_start_date' in honeypots.keys():
             honeypots['operation_start_date'] = pd.to_datetime(honeypots['operation_start_date'])
+
+        # Parsing end date, filling empty values with date of 'today'
         if 'operation_end_date' in honeypots.keys():
             honeypots['operation_end_date'] = honeypots['operation_end_date'].fillna(dt.date.today())
             honeypots['operation_end_date'] = pd.to_datetime(honeypots['operation_end_date'])
+
+        # Keep honeypots active on the specified date
         if ('operation_start_date' in honeypots.keys()) and 'operation_end_date' in honeypots.keys():
             honeypots = honeypots[(for_date >= honeypots['operation_start_date']) & (for_date <= honeypots['operation_end_date'])]
+
     ips = honeypots.public_ip.values
     return ips
 
-def _process_zeek_files(zeek_files, date):
+
+def _process_zeek_files(list_of_zeek_files, date):
+    """
+    Process a list of Zeek log files to extract all connections
+    sent to the honeypot IPs (the attacks) for a given date.
+    """
+    # Retrieve the list of honeypot IPs
     ips = _get_honeypot_ips()
+
+    # Initialises daily, a dataframe that will contain
+    # all the connections to the honeypot IPs found
+    # on the input zeek files
     daily = pd.DataFrame()
-    for z in zeek_files:
+
+    # Process each zeek file in the input list
+    for zeek_file in list_of_zeek_files:
         hourly = pd.DataFrame()
-        zeekdata = read_zeek(z)
+
+        # Read the zeek file into a dataframe
+        zeekdata = read_zeek(zeek_file)
+
+        # Keep the connections where the honeypot is the responder,
+        # that is, traffic sent to the honeypot
         for ip in ips:
             hourly = pd.concat([hourly, zeekdata[zeekdata['id.resp_h'] == ip]])
+
+        # Store the hourly traffic on the daily dataframe
         daily = pd.concat([daily, hourly])
+
+    # Return a DF with all the traffic sent to the honeypot IPs
+    # on the input Zeek files
     return daily
 
+
 def _process_argus_files(argus_files, date):
     ips = _get_honeypot_ips()
     daily = pd.DataFrame()
@@ -93,6 +135,7 @@ def _process_argus_files(argus_files, date):
         daily = pd.concat([daily, hourly])
     return daily
 
+
 def _process_raw_files(date):
     '''
     Create a dataset for the date string date in the data/interim folder
@@ -102,7 +145,7 @@ def _process_raw_files(date):
     # if data directory does not exist, execute the magic to get it
     if path.isdir(path.join(project_dir,'data','raw', date)) == False:
         logging.debug(f'Downloading data for {date}')
-        getrawdata(date)
+        get_raw_data(date)
     # after this point, if directory does not exist, we can skip it.
     try:
         zeek_files = [x.path for x in scandir(path.join(project_dir,'data','raw', date)) if x.name.startswith('conn.')]
@@ -116,9 +159,10 @@ def _process_raw_files(date):
     daily.to_csv(path.join(project_dir,'data','interim', f'daily.conn.{date}.csv.gz'), index=False, compression='gzip')
     logger.debug('Writting file: ' + path.join(project_dir,'data','interim', f'daily.conn.{date}.csv.gz'))
     #logger.debug('Removing raw data (not needed anymore): ' + path.join(project_dir,'data','raw', f'{date}'))
-    #removerawdata(date)
+    #remove_raw_data(date)
     return
 
+
 def _extract_attacks(date):
     '''
     Create a dataset for the date string date in the data/interim folder
@@ -150,9 +194,10 @@ def _extract_attacks(date):
     logger.debug('Writting file: ' + path.join(project_dir,'data','processed', f'attacks.{date}.csv.gz'))
     df.to_csv(path.join(project_dir,'data','processed', f'attacks.{date}.csv.gz'), columns=['orig', 'flows', 'duration', 'packets', 'bytes'], index=False, compression='gzip')
     # logger.debug('Removing raw data (not needed anymore): ' + path.join(project_dir,'data','raw', f'{date}'))
-    #removerawdata(date)
+    #remove_raw_data(date)
     return
 
+
 def process_zeek_files(dates=None):
     """ 
     Creates the dataset or part of it
@@ -170,6 +215,7 @@ def process_zeek_files(dates=None):
     Parallel(n_jobs=12, backend='multiprocessing')(delayed(_process_raw_files)(date) for date in dates)
     return
 
+
 def extract_attacks(dates=None):
     """
     Creates the dataset or part of it
@@ -193,6 +239,7 @@ def extract_attacks(dates=None):
     Parallel(n_jobs=12, backend='multiprocessing')(delayed(_extract_attacks)(date) for date in dates)
     return
 
+
 def get_attacks(start=None, end=None, dates=None, usecols=None):
     '''
     Returns a DataFrame with the attacks between the dates start and end or the
@@ -212,6 +259,7 @@ def get_attacks(start=None, end=None, dates=None, usecols=None):
             for date in dates]
     return dfs
 
+
 if __name__ == '__main__':
     log_fmt = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
     #logging.basicConfig(level=logging.INFO, format=log_fmt)

lib/aip/data/functions.py

@@ -3,28 +3,44 @@
 import hashlib
 import pandas as pd
 import shutil
-import subprocess, shlex
+import shlex
+import subprocess
 import zeeklog2pandas as z2p
 
 from dotenv import dotenv_values
-from joblib import Parallel, delayed
-from os import makedirs, path, access, W_OK
+from joblib import Parallel
+from joblib import delayed
+from os import makedirs
+from os import path
+from os import access
+from os import W_OK
 from pathlib import Path
 
 _project_dir = Path(__file__).resolve().parents[3]
 _config = {
     **dotenv_values(path.join(_project_dir, ".env")),  # load sensitive variables
 }
 
+
 def read_zeek(path, **kwargs):
+    """
+    Reads a Zeek file to a DataFrame
+    """
     try:
+        # Load Zeek path to DataFrame
         df = z2p.read_zeek(path, **kwargs)
+
+        # Convert to readable time format
         if 'ts' in df.keys():
             df['ts'] = pd.to_datetime(df.ts, unit='s')
+
+        # Returns data frame
         return df
     except:
         raise z2p.NotAZeekFile(path)
 
+
+# Deprecated, AIP only reads Zeek logs
 def read_argus(path, **kwargs):
     from os import path as ospath
     if ospath.exists(path.path + '.csv'):
@@ -40,21 +56,55 @@ def read_argus(path, **kwargs):
     df.rename(columns={'StartTime':'ts', 'SrcAddr':'id.orig_h', 'DstAddr':'id.resp_h', 'Dur': 'duration', 'SrcBytes': 'orig_ip_bytes', 'SrcPkts': 'orig_pkts'}, inplace=True)
     return df
 
+
+# This function is unused right now
 def scramble(s):
     return hashlib.sha1(_config['salt'].encode() + s.encode()).hexdigest()
 
-def getrawdata(date):
+
+def get_raw_data(date):
+    """
+    Copies the Zeek conn logs of a given date into data/raw/YYYY-MM-DD,
+    using the copy command set in 'magic' in the .env file (for example
+    'cp' or 'scp'). The 24 hourly files are copied in parallel.
+
+    Known issue: this function never copies anything. It is only called
+    when data/raw/YYYY-MM-DD does not exist, and access() returns False
+    for a directory that does not exist, so the copy is always skipped.
+    """
+
+    # Validate date is well formatted
     dt.datetime.strptime(date, '%Y-%m-%d')
-    p = path.join(_project_dir,'data','raw', date)
-    if access(p, W_OK):
-        makedirs(p, exist_ok=True)
-        commands = [shlex.split(_config['magic'] + f'{date}/conn.{x:02}* ' + p) for x in range(0,24)]
-        Parallel(n_jobs=24, backend='threading')(delayed(subprocess.run)(c) for c in commands)
 
-def removerawdata(date, force=False):
+    raw_data_dir = path.join(_project_dir,'data','raw', date)
+
+    # Only copy if the directory exists and is writable (see known issue above)
+    if access(raw_data_dir, W_OK):
+        # Create directory, ignore if it exists
+        makedirs(raw_data_dir, exist_ok=True)
+
+        # Build one copy command per hour of the day:
+        # <magic><date>/conn.<hour>* data/raw/YYYY-MM-DD
+        commands = [
+            shlex.split(_config['magic'] + f'{date}/conn.{x:02}* ' + raw_data_dir)
+            for x in range(0,24)
+        ]
+
+        # Run the copy commands in parallel
+        Parallel(n_jobs=24, backend='threading')(delayed(subprocess.run)(cmd) for cmd in commands)
+
+
+def remove_raw_data(date, force=False):
+    """
+    Delete the raw data directory of a given date. Only runs if
+    remove_raw_data is set to true in the .env file, or if force=True.
+    """
+    # Validate date is well formatted
     dt.datetime.strptime(date, '%Y-%m-%d')
-    p = path.join(_project_dir,'data','raw', date)
+
+    raw_data_dir = path.join(_project_dir,'data','raw', date)
+
     # Only delete raw data if explicitly allowed in the configuration file
     if (_config['remove_raw_data'].lower() == 'true') or force:
-        shutil.rmtree(p)
+        shutil.rmtree(raw_data_dir)