From cfd8ccb1e7557870a60e509ea12132cae47132cb Mon Sep 17 00:00:00 2001 From: Veronica Valeros Date: Mon, 28 Oct 2024 16:52:21 +0100 Subject: [PATCH] Add argparse to bin/aip (#51) * Refactor imports * Fix bad date reference (introduced error) * Add argparse with date and model options * Run if model selected or all * Move date funct. to utils * Create output dir on run_model to simplify logic * Add param to adjust logging level --- bin/aip | 71 +++++++++++++++++++------------------ lib/aip/utils/date_utils.py | 14 ++++++++ 2 files changed, 51 insertions(+), 34 deletions(-) create mode 100644 lib/aip/utils/date_utils.py diff --git a/bin/aip b/bin/aip index 3b712b5..d28055b 100755 --- a/bin/aip +++ b/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() diff --git a/lib/aip/utils/date_utils.py b/lib/aip/utils/date_utils.py new file mode 100644 index 0000000..6cdf829 --- /dev/null +++ b/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 \ No newline at end of file