Add help and arguments to bin/aip

bin/aip

@@ -29,69 +29,72 @@ __license__ = "GPLv3"
 __maintainer__ = "Joaquin Bogado"
 __version__ = "1.0.0"
 
+import argparse
 import logging
-import pandas as pd
-
-from aip.data.access import data_path, project_dir
+from datetime import date
+from os import makedirs
+from os import path
+from aip.data.access import data_path
 from aip.models.alpha import Alpha
 from aip.models.prioritize import New
 from aip.models.prioritize import Consistent
 from aip.models.prioritize import RandomForest
-from os import makedirs, path, scandir
-from datetime import date, timedelta, datetime
-import sys
-
+from aip.utils.date_utils import validate_and_convert_date
 
-def validate_and_convert_date(date_str):
-    """
-    Validates a date string in 'YYYY-MM-DD' format and converts it to a date object.
-    """
-    try:
-        dateobj = datetime.strptime(date_str, '%Y-%m-%d')
-        return dateobj.date()
-    except ValueError as e:
-        print('Invalid date format. It should be YYYY-MM-DD')
-        raise e
 
-def run_model(aip_model_name, aip_output_dir, aip_model, date_day):
+def run_model(aip_model_name, aip_model, date_day):
     """
     Run a given model with exception handling
     """
     blocklist=""
+    model_output_dir = path.join(data_path,'output',aip_model_name)
     # Make sure output directory is created
-    if not path.exists(aip_output_dir):
-        makedirs(aip_output_dir)
+    if not path.exists(model_output_dir):
+        makedirs(model_output_dir)
 
     try:
-        blocklist = aip_model.run(day)
-        blocklist.to_csv(path.join(aip_output_dir, f'AIP-{aip_model_name}-{str(date_day)}.csv.gz'), index=False, compression='gzip')
+        blocklist = aip_model.run(date_day)
+        blocklist.to_csv(path.join(model_output_dir, f'AIP-{aip_model_name}-{str(date_day)}.csv.gz'), index=False, compression='gzip')
         logging.info(f"{aip_model_name} model completed successfully.")
     except Exception as e:
         logging.error(f"Error running {aip_model_name} model: {e}", exc_info=True)
 
 
-if __name__ == '__main__':
-    if len(sys.argv) == 2:
-        datestr = sys.argv[1]
-        day = validate_and_convert_date(datestr)
-    else:
-        day = date.today()
+def main():
+    parser = argparse.ArgumentParser(description='Attacker IP Prioritization (AIP) Tool')
+    parser.add_argument('--date', type=str, help='The date for running the models in YYYY-MM-DD format. Defaults to today.', default=str(date.today()))
+    parser.add_argument('--model', type=str, choices=['Alpha', 'Alpha7', 'Prioritize_New', 'Prioritize_Consistent', 'Random_Forest', 'all'], default='all', help='Select AIP model to run. Defaults to all.')
+    parser.add_argument('-d', '--debug', required=False, help="Debugging mode.", action="store_const", dest="log_level", const=logging.DEBUG, default=logging.ERROR,)
+    parser.add_argument('-v', '--verbose', required=False, help="Verbose mode", action="store_const", dest="log_level", const=logging.INFO,)
+
+    args = parser.parse_args()
 
     # Set up logging
     log_fmt = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-    logging.basicConfig(level=logging.DEBUG, format=log_fmt)
+    logging.basicConfig(level=args.log_level, format=log_fmt)
+
+    # Validate input date
+    run_date_day = validate_and_convert_date(args.date)
 
     # Run Alpha Model
-    run_model('Alpha', path.join(data_path, 'output', 'Alpha'), Alpha(), day)
+    if args.model in ['Alpha', 'all']:
+        run_model('Alpha', Alpha(), run_date_day)
 
     # Alpha 7 Model
-    run_model('Alpha7', path.join(data_path, 'output', 'Alpha7'), Alpha(lookback=7), day)
+    if args.model in ['Alpha7', 'all']:
+        run_model('Alpha7', Alpha(lookback=7), run_date_day)
 
     # Prioritize New Model
-    run_model('Prioritize_New', path.join(data_path, 'output', 'Prioritize_New'), New(), day)
+    if args.model in ['Prioritize_New', 'all']:
+        run_model('Prioritize_New', New(), run_date_day)
 
     # Prioritize Consistent Model
-    run_model('Prioritize_Consistent', path.join(data_path, 'output', 'Prioritize_Consistent'), Consistent(), day)
+    if args.model in ['Prioritize_Consistent', 'all']:
+        run_model('Prioritize_Consistent', Consistent(), run_date_day)
 
     # Prioritize Random Forest Model
-    run_model('Random_Forest', path.join(data_path, 'output', 'random_forest'), RandomForest(), day)
+    if args.model in ['Random_Forest', 'all']:
+        run_model('Random_Forest', RandomForest(), run_date_day)
+
+if __name__ == '__main__':
+    main()

lib/aip/utils/date_utils.py

@@ -0,0 +1,14 @@
+import logging
+from datetime import datetime
+
+
+def validate_and_convert_date(date_str):
+    """
+    Validates a date string in 'YYYY-MM-DD' format and converts it to a date object.
+    """
+    try:
+        dateobj = datetime.strptime(date_str, '%Y-%m-%d')
+        return dateobj.date()
+    except ValueError as e:
+        logging.error(f"Invalid date format for '{date_str}', expected YYYY-MM-DD")
+        raise ValueError(f"Invalid date format: {date_str}, expected YYYY-MM-DD") from e