Add scalable AIP arguments

README.md

@@ -53,12 +53,6 @@ AIP will automatically attempt to run all the models using the available data. A
 :~$ docker run --rm -v /opt/zeek/logs/:/home/aip/AIP/data/raw:ro -v ${PWD}/data/:/home/aip/AIP/data/:rw --name aip stratosphereips/aip:latest bin/aip
 ```
 
-To run AIP for a specific day:
-```bash
-:~$ cd AIP
-:~$ docker run --rm -v /opt/zeek/logs/:/home/aip/AIP/data/raw:ro -v ${PWD}/data/:/home/aip/AIP/data/:rw --name aip stratosphereips/aip:latest bin/aip YYYY-MM-DD
-```
-
 ### Local Installation
 
 To run AIP locally, follow the next steps:
@@ -78,6 +72,17 @@ To run AIP locally, follow the next steps:
 * Let's test that AIP works with a simple test:
     * `python3 bin/aip --help`
 
+### AIP Commands
+
+Get all the full list of options with `bin/aip --help`:
+
+- Run AIP for a specific day:
+    - `bin/aip -v run-models --start-date YYYY-MM-DD`
+- Run AIP for a specific day and model:
+    - `bin/aip -v run-models --start-date YYYY-MM-DD --model Alpha`
+- Rebuild AIP knowledge base from a specific date:
+    - `bin/aip -d rebuild-kb --start-date YYYY-MM-DD`
+
 ## License
 
 The Stratosphere AIP tool is licensed under [GNU General Public License v3.0](https://github.com/stratosphereips/AIP/blob/main/LICENSE).

bin/aip

@@ -67,49 +67,64 @@ def main():
     parser = argparse.ArgumentParser(description='Attacker IP Prioritization (AIP) Tool')
     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,)
-    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('--rebuild-kb', type=str, help='Rebuild knowledge base from a specific date (YYYY-MM-DD).')
+
+    subparsers = parser.add_subparsers(dest="command", help="commands")
+
+    # Run Models command 
+    run_models_parser = subparsers.add_parser("run-models", help="Run models with optional date range.")
+    run_models_parser.add_argument('--start-date', type=str, help='The start date for running the models (YYYY-MM-DD). Default is today.', default=str(date.today()))
+    run_models_parser.add_argument('--model', type=str, choices=['Alpha', 'Alpha7', 'Prioritize_New', 'Prioritize_Consistent', 'Random_Forest', 'all'], default='all', help='Select AIP model to run. Default is all models.')
+
+    # Rebuild KB command
+    rebuild_kb_parser = subparsers.add_parser("rebuild-kb", help="Rebuild the knowledge base from a specific date.")
+    rebuild_kb_parser.add_argument("--start-date", required=True, type=str, help="The date from which to rebuild the knowledge base in YYYY-MM-DD format.")
 
     args = parser.parse_args()
 
     # Set up logging
-    logger = logging.getLogger(__name__)
+    logger = logging.getLogger('aip')
     log_fmt = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
     logging.basicConfig(level=args.log_level, format=log_fmt)
 
     try: 
+        # Validate start_date once not on every command
+        if hasattr(args,'start_date'):
+            aip_start_date = validate_and_convert_date(args.start_date)
+        else:
+            # If not set, default to today()
+            aip_start_date = date.today()
+
+        # If model is not set, default to today
+        if not hasattr(args,'model'):
+            args.model = 'all'
+
         # Rebuild knowledge base
-        if args.rebuild_kb:
-            rebuild_date = validate_and_convert_date(args.rebuild_kb)
-            logger.info(f"Rebuilding knowledge base from {rebuild_date} onwards")
-            _rebuild(start_date=rebuild_date, log_level=args.log_level)
+        if args.command == "rebuild-kb":
+            logger.info(f"Rebuilding knowledge base from {aip_start_date} onwards")
+            _rebuild(start_date=aip_start_date, log_level=args.log_level)
             sys.exit()
 
-        # Validate input date
-        if args.date: 
-            run_date_day = validate_and_convert_date(args.date)
-
-        logging.info(f"Running {args.model} model(s) for date {run_date_day}.")
+        logging.info(f"Running {args.model} model(s) for date {aip_start_date}.")
         # Run Alpha Model
         if args.model in ['Alpha', 'all']:
-            run_model('Alpha', Alpha(), run_date_day, args.log_level)
+            run_model('Alpha', Alpha(), aip_start_date, args.log_level)
 
         # Alpha 7 Model
         if args.model in ['Alpha7', 'all']:
-            run_model('Alpha7', Alpha7(), run_date_day, args.log_level)
+            run_model('Alpha7', Alpha7(), aip_start_date, args.log_level)
 
         # Prioritize New Model
         if args.model in ['Prioritize_New', 'all']:
-            run_model('Prioritize_New', New(), run_date_day, args.log_level)
+            run_model('Prioritize_New', New(), aip_start_date, args.log_level)
 
         # Prioritize Consistent Model
         if args.model in ['Prioritize_Consistent', 'all']:
-            run_model('Prioritize_Consistent', Consistent(), run_date_day, args.log_level)
+            run_model('Prioritize_Consistent', Consistent(), aip_start_date, args.log_level)
 
         # Prioritize Random Forest Model
         if args.model in ['Random_Forest', 'all']:
-            run_model('Random_Forest', RandomForest(), run_date_day, args.log_level)
+            run_model('Random_Forest', RandomForest(), aip_start_date, args.log_level)
+
     except ValueError as err:
         logger.error(err)
         sys.exit()