From 5f2ed3d233edb45387c4bdeaa903737e8758d14c Mon Sep 17 00:00:00 2001 From: verovaleros Date: Tue, 5 Nov 2024 17:10:45 +0000 Subject: [PATCH 1/8] Fix logger name to smth meaningful --- bin/aip | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bin/aip b/bin/aip index f3879b9..484f29e 100755 --- a/bin/aip +++ b/bin/aip @@ -74,7 +74,7 @@ def main(): 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) From 0cbdfd49bf30c9e88c7c228ddd9312383af79ade Mon Sep 17 00:00:00 2001 From: verovaleros Date: Tue, 5 Nov 2024 17:11:45 +0000 Subject: [PATCH 2/8] Add subparser --- bin/aip | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/bin/aip b/bin/aip index 484f29e..7246d8d 100755 --- a/bin/aip +++ b/bin/aip @@ -67,9 +67,9 @@ 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") + args = parser.parse_args() From 7407b4bd4cc7dee169035fb8d1289b68949c58aa Mon Sep 17 00:00:00 2001 From: verovaleros Date: Tue, 5 Nov 2024 17:12:06 +0000 Subject: [PATCH 3/8] Migrate to 'run-models' command --- bin/aip | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/bin/aip b/bin/aip index 7246d8d..df1e464 100755 --- a/bin/aip +++ b/bin/aip @@ -70,6 +70,12 @@ def main(): 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 args = parser.parse_args() From bfec0ad3b9c247dc62c0240b90ba997f2eed7a44 Mon Sep 17 00:00:00 2001 From: verovaleros Date: Tue, 5 Nov 2024 17:12:23 +0000 Subject: [PATCH 4/8] Migrate to 'rebuild-kb' command --- bin/aip | 2 ++ 1 file changed, 2 insertions(+) diff --git a/bin/aip b/bin/aip index df1e464..1e92868 100755 --- a/bin/aip +++ b/bin/aip @@ -76,6 +76,8 @@ def main(): 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() From 2ef6844f0a0f8b92bda6a8c586360d44e2625ed5 Mon Sep 17 00:00:00 2001 From: verovaleros Date: Tue, 5 Nov 2024 17:12:47 +0000 Subject: [PATCH 5/8] Validate start_date once, not on every command --- bin/aip | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/bin/aip b/bin/aip index 1e92868..cc85f48 100755 --- a/bin/aip +++ b/bin/aip @@ -87,6 +87,12 @@ def main(): 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() # Rebuild knowledge base if args.rebuild_kb: rebuild_date = validate_and_convert_date(args.rebuild_kb) From 7c2b6cfcc4f222efe3dc2f60b6a5e8eef7e93ac0 Mon Sep 17 00:00:00 2001 From: verovaleros Date: Tue, 5 Nov 2024 17:13:02 +0000 Subject: [PATCH 6/8] Set models default when not specified --- bin/aip | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/bin/aip b/bin/aip index cc85f48..241e406 100755 --- a/bin/aip +++ b/bin/aip @@ -93,6 +93,11 @@ def main(): 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) From 196bfa8172af1d31c3a9576a6c6bbc25e394a05f Mon Sep 17 00:00:00 2001 From: verovaleros Date: Tue, 5 Nov 2024 17:13:19 +0000 Subject: [PATCH 7/8] Update to reflect change in variable names --- bin/aip | 24 ++++++++++-------------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/bin/aip b/bin/aip index 241e406..c431564 100755 --- a/bin/aip +++ b/bin/aip @@ -99,36 +99,32 @@ def main(): 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() From eb8e59ba48488d7598393a0fee7701a6250d9157 Mon Sep 17 00:00:00 2001 From: verovaleros Date: Tue, 5 Nov 2024 17:53:38 +0000 Subject: [PATCH 8/8] Add some commands to the README --- README.md | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 18f3a3e..d00bf2d 100644 --- a/README.md +++ b/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).