diff --git a/dettect.py b/dettect.py index 8f5c424..7fe4ff8 100644 --- a/dettect.py +++ b/dettect.py @@ -131,9 +131,12 @@ def _init_menu(): help='includes: statistics on ATT&CK data source and updates on techniques' ', groups and software', aliases=['ge']) - parser_generic.add_argument('-s', '--statistics', help='get a sorted count on how much techniques are covered by a ' - 'particular data source', action='store_true') - + parser_generic.add_argument('-ds', '--datasources', help='get a sorted count on how many ATT&CK Enterprise ' + 'techniques are covered by a particular Data Source', + action='store_true') + parser_generic.add_argument('-m', '--mitigations', help='get a sorted count on how many ATT&CK Enterprise or ' + 'Mobile techniques are covered by a Mitigation', + choices=['enterprise', 'mobile']) parser_generic.add_argument('-u', '--updates', help='get a sorted list for when updates were released for ' 'techniques, groups or software', choices=['techniques', 'groups', 'software']) @@ -211,8 +214,10 @@ def _menu(menu_parser): print("[!] Filtering on 'applicable_to' is not supported for Excel output") elif args.subparser in ['generic', 'ge']: - if args.statistics: - get_statistics() + if args.datasources: + get_statistics_data_sources() + elif args.mitigations: + get_statistics_mitigations(args.mitigations) elif args.updates: get_updates(args.updates, args.sort) diff --git a/interactive_menu.py b/interactive_menu.py index 63b3909..75c3d2b 100644 --- a/interactive_menu.py +++ b/interactive_menu.py @@ -7,8 +7,9 @@ from constants import * groups = 'all' software_group = False -platform = 'Windows' -stage = 'attack' +default_platform = 'Windows' +default_stage = 'attack' +default_matrix = 'enterprise' groups_overlay = '' overlay_type = 'group' filter_applicable_to = 'all' @@ -194,12 +195,31 @@ def _menu_statistics(): Handles the Statistics functionality. :return: """ + global default_matrix _clear() print('Menu: Statistics') print('') - get_statistics() + print('Options:') + print('1. Matrix: %s' % default_matrix) + print('2. Get a sorted count on how many ATT&CK Enterprise techniques are covered by a particular Data Source.') + print('3. Get a sorted count on how many ATT&CK Enterprise or Mobile techniques are covered by a Mitigation.') + print('9. Back to main menu.') + choice = _ask_input() + if choice == '1': + print('Specify the matrix (enterprise or mobile):') + m = _ask_input().lower() + default_matrix = 'enterprise' if m == 'enterprise' else 'mobile' + elif choice == '2': + get_statistics_data_sources() + elif choice == '3': + get_statistics_mitigations(default_matrix) + elif choice == '9': + interactive_menu() + elif choice == 'q': + quit() + _wait() - interactive_menu() + _menu_statistics() def _menu_data_source(filename_ds): @@ -366,14 +386,14 @@ def _menu_groups(): Prints and handles the Threat actor group mapping functionality. :return: """ - global groups, software_group, platform, stage, groups_overlay, overlay_type, filter_applicable_to + global groups, software_group, default_platform, default_stage, groups_overlay, overlay_type, filter_applicable_to _clear() print('Menu: %s' % MENU_NAME_THREAT_ACTOR_GROUP_MAPPING) print('') print('Options:') print('1. Software group: %s' % str(software_group)) - print('2. Platform: %s' % platform) - print('3. Stage: %s' % stage) + print('2. Platform: %s' % default_platform) + print('3. Stage: %s' % default_stage) print('4. Groups: %s' % groups) print('5. Overlay: ') print(' - %s: %s' % ('File' if os.path.exists(groups_overlay) else 'Groups', groups_overlay)) @@ -390,11 +410,11 @@ def _menu_groups(): elif choice == '2': print('Specify platform (all, Linux, macOS, Windows):') p = _ask_input().lower() - platform = 'Windows' if p == 'windows' else 'Linux' if p == 'linux' else 'macOS' if p == 'macos' else 'all' + default_platform = 'Windows' if p == 'windows' else 'Linux' if p == 'linux' else 'macOS' if p == 'macos' else 'all' elif choice == '3': print('Specify stage (pre-attack, attack):') s = _ask_input().lower() - stage = 'pre-attack' if s == 'pre-attack' else 'attack' + default_stage = 'pre-attack' if s == 'pre-attack' else 'attack' elif choice == '4': print('Specify the groups to include separated using commas. Group can be their ID, name or alias ' '(default is all groups). Other option is to provide a YAML file with a custom group(s)') @@ -426,7 +446,7 @@ def _menu_groups(): print('Specify your filter for the applicable_to field:') filter_applicable_to = _ask_input().lower() elif choice == '7': - generate_group_heat_map(groups, groups_overlay, overlay_type, stage, platform, software_group, filter_applicable_to) + generate_group_heat_map(groups, groups_overlay, overlay_type, default_stage, default_platform, software_group, filter_applicable_to) _wait() elif choice == '9': interactive_menu()