- Multiple functions made "private".

- Added new menu options.
master
Marcus Bakker 2019-07-31 10:20:21 +02:00
parent f4121bf4d0
commit b2fdb25647
1 changed files with 95 additions and 87 deletions

View File

@ -1,4 +1,3 @@
import sys
import glob
from data_source_mapping import *
from technique_mapping import *
@ -16,7 +15,7 @@ filter_applicable_to = 'all'
yaml_path = 'sample-data/'
def clear():
def _clear():
"""
Clears the terminal screen and prints the title and version of the application.
:return:
@ -34,7 +33,7 @@ def clear():
print('')
def ask_input():
def _ask_input():
"""
Waits for input from the terminal.
:return:
@ -42,7 +41,7 @@ def ask_input():
return input(' >> ')
def wait():
def _wait():
"""
Prints wait statement and wait for pressing ENTER key.
:return:
@ -57,7 +56,7 @@ def interactive_menu():
Main menu for interactive mode.
:return:
"""
clear()
_clear()
print('Select a mode:')
print('1. %s' % MENU_NAME_DATA_SOURCE_MAPPING)
print('2. %s' % MENU_NAME_VISIBILITY_MAPPING)
@ -66,38 +65,38 @@ def interactive_menu():
print('5. Updates')
print('6. Statistics')
print('9. Quit')
choice = ask_input()
choice = _ask_input()
if choice == '1':
menu_data_source(select_file(MENU_NAME_DATA_SOURCE_MAPPING, 'data sources', FILE_TYPE_DATA_SOURCE_ADMINISTRATION))
_menu_data_source(_select_file(MENU_NAME_DATA_SOURCE_MAPPING, 'data sources', FILE_TYPE_DATA_SOURCE_ADMINISTRATION))
elif choice == '2':
menu_visibility(select_file(MENU_NAME_VISIBILITY_MAPPING, 'techniques (used to score the level of visibility)', FILE_TYPE_TECHNIQUE_ADMINISTRATION),
select_file(MENU_NAME_VISIBILITY_MAPPING, 'data sources (used to add metadata on the involved data sources to the heat map)', FILE_TYPE_DATA_SOURCE_ADMINISTRATION, False))
_menu_visibility(_select_file(MENU_NAME_VISIBILITY_MAPPING, 'techniques (used to score the level of visibility)', FILE_TYPE_TECHNIQUE_ADMINISTRATION),
_select_file(MENU_NAME_VISIBILITY_MAPPING, 'data sources (used to add metadata on the involved data sources to the heat map)', FILE_TYPE_DATA_SOURCE_ADMINISTRATION, False))
elif choice == '3':
menu_detection(select_file(MENU_NAME_DETECTION_COVERAGE_MAPPING, 'techniques', FILE_TYPE_TECHNIQUE_ADMINISTRATION))
_menu_detection(_select_file(MENU_NAME_DETECTION_COVERAGE_MAPPING, 'techniques', FILE_TYPE_TECHNIQUE_ADMINISTRATION))
elif choice == '4':
menu_groups()
_menu_groups()
elif choice == '5':
menu_updates()
_menu_updates()
elif choice == '6':
menu_statistics()
_menu_statistics()
elif choice in ['9', 'q']:
quit()
else:
interactive_menu()
def select_file(title, what, expected_file_type, b_clear=True):
def _select_file(title, what, expected_file_type, b_clear=True):
"""
Prints and handles the file selection in the terminal. It shows just .yaml files.
:param title: title to print on top of this menu
:param what: print for what purpose the file is selected
:param expected_file_type: the expected file type of the YAML file
:param b_clear: clear the terminal before showing this menu
:param b_clear: _clear the terminal before showing this menu
:return: filename of the selected file
"""
global yaml_path
if b_clear:
clear()
_clear()
print('Menu: %s' % title)
print('')
print('Select the YAML file with %s:' % what)
@ -116,18 +115,18 @@ def select_file(title, what, expected_file_type, b_clear=True):
back_nr = 9 if n < 9 else n + (5 - n % 5)
print('%d. Back to main menu.' % back_nr)
choice = ask_input()
choice = _ask_input()
if choice == str(change_path_nr):
print("Supply full or relative path:")
choice = ask_input()
choice = _ask_input()
choice = choice if choice.endswith('/') else choice + '/'
if os.path.exists(choice):
yaml_path = choice
return select_file(title, what, expected_file_type, b_clear)
return _select_file(title, what, expected_file_type, b_clear)
else:
print("[!] Path doesn't exist")
wait()
return select_file(title, what, expected_file_type, b_clear)
_wait()
return _select_file(title, what, expected_file_type, b_clear)
elif choice == str(back_nr):
interactive_menu()
elif choice == 'q':
@ -138,21 +137,21 @@ def select_file(title, what, expected_file_type, b_clear=True):
file_type = check_file(filename, file_type=expected_file_type)
if file_type:
print('Selected file: ' + filename)
wait()
_wait()
return filename
else:
print("[!] Invalid choice")
wait()
return select_file(title, what, expected_file_type, b_clear)
_wait()
return _select_file(title, what, expected_file_type, b_clear)
def menu_updates():
def _menu_updates():
"""
Prints and handles the menu for the Updates functionality.
:return:
"""
clear()
_clear()
print('Menu: Updates')
print('')
@ -164,55 +163,55 @@ def menu_updates():
print('3. Software (sorted by modified date)')
print('3s. Software (sorted by creation date)')
print('9. Back to main menu.')
choice = ask_input()
choice = _ask_input()
if choice == '1':
get_updates('techniques')
wait()
_wait()
if choice == '1s':
get_updates('techniques', 'created')
wait()
_wait()
elif choice == '2':
get_updates('groups')
wait()
_wait()
elif choice == '2s':
get_updates('groups', 'created')
wait()
_wait()
elif choice == '3':
get_updates('software')
wait()
_wait()
elif choice == '3s':
get_updates('software', 'created')
wait()
_wait()
elif choice == '9':
interactive_menu()
elif choice == 'q':
quit()
menu_updates()
_menu_updates()
def menu_statistics():
def _menu_statistics():
"""
Handles the Statistics functionality.
:return:
"""
clear()
_clear()
print('Menu: Statistics')
print('')
get_statistics()
wait()
_wait()
interactive_menu()
def menu_data_source(filename):
def _menu_data_source(filename_ds):
"""
Prints and handles the Data source mapping functionality.
:param filename:
:param filename_ds:
:return:
"""
clear()
_clear()
print('Menu: %s' % MENU_NAME_DATA_SOURCE_MAPPING)
print('')
print('Selected data source YAML file: %s' % filename)
print('Selected data source YAML file: %s' % filename_ds)
print('')
print('Select what you want to do:')
print('1. Generate a data source layer for the ATT&CK Navigator.')
@ -220,39 +219,48 @@ def menu_data_source(filename):
print('3. Generate an Excel sheet with all data sources.')
print('4. Generate a technique administration YAML file with visibility scores, based on the number of available '
'data sources')
print('5. update the visibility scores within a technique administration YAML file based on changes within any of '
'the data sources. \nPast visibility scores are preserved in the score_logbook, and manually assigned scores are '
'not updated without your approval. \nThe updated visibility are based on the number of available data sources.')
print('9. Back to main menu.')
choice = ask_input()
choice = _ask_input()
if choice == '1':
print('Writing data sources layer...')
generate_data_sources_layer(filename)
wait()
generate_data_sources_layer(filename_ds)
_wait()
elif choice == '2':
print('Drawing the graph...')
plot_data_sources_graph(filename)
wait()
plot_data_sources_graph(filename_ds)
_wait()
elif choice == '3':
print('Generating Excel file...')
export_data_source_list_to_excel(filename)
wait()
export_data_source_list_to_excel(filename_ds)
_wait()
elif choice == '4':
print('Generating YAML file...')
generate_technique_administration_file(filename)
wait()
generate_technique_administration_file(filename_ds)
_wait()
elif choice == '5':
filename_t = _select_file(MENU_NAME_DETECTION_COVERAGE_MAPPING, 'techniques (used to score the level of visibility)',
FILE_TYPE_TECHNIQUE_ADMINISTRATION, False)
print('Updating visibility scores...')
update_technique_administration_file(filename_ds, filename_t)
_wait()
elif choice == '9':
interactive_menu()
elif choice == 'q':
quit()
menu_data_source(filename)
_menu_data_source(filename_ds)
def menu_detection(filename_t):
def _menu_detection(filename_t):
"""
Prints and handles the Detection coverage mapping functionality.
:param filename_t:
:return:
"""
global filter_applicable_to
clear()
_clear()
print('Menu: %s' % MENU_NAME_DETECTION_COVERAGE_MAPPING)
print('')
print('Selected techniques YAML file: %s' % filename_t)
@ -268,41 +276,41 @@ def menu_detection(filename_t):
print('5. Generate an Excel sheet with all administrated techniques.')
print('6. Check the technique YAML file for errors.')
print('9. Back to main menu.')
choice = ask_input()
choice = _ask_input()
if choice == '1':
print('Specify your filter for the applicable_to field:')
filter_applicable_to = ask_input().lower()
filter_applicable_to = _ask_input().lower()
elif choice == '2':
print('Writing detection coverage layer...')
generate_detection_layer(filename_t, None, False, filter_applicable_to)
wait()
_wait()
elif choice == '3':
filename_ds = select_file(MENU_NAME_DETECTION_COVERAGE_MAPPING, 'data sources (used to add metadata on the '
filename_ds = _select_file(MENU_NAME_DETECTION_COVERAGE_MAPPING, 'data sources (used to add metadata on the '
'involved data sources to the heat map)',
FILE_TYPE_DATA_SOURCE_ADMINISTRATION, False)
print('Writing detection coverage layer with visibility as overlay...')
generate_detection_layer(filename_t, filename_ds, True, filter_applicable_to)
wait()
_wait()
elif choice == '4':
print('Drawing the graph...')
plot_detection_graph(filename_t, filter_applicable_to)
wait()
_wait()
elif choice == '5':
print('Generating Excel file...')
export_techniques_list_to_excel(filename_t)
wait()
_wait()
elif choice == '6':
print('Checking the technique YAML file for errors...')
check_yaml_file_health(filename_t, FILE_TYPE_TECHNIQUE_ADMINISTRATION, health_is_called=True)
wait()
_wait()
elif choice == '9':
interactive_menu()
elif choice == 'q':
quit()
menu_detection(filename_t)
_menu_detection(filename_t)
def menu_visibility(filename_t, filename_ds):
def _menu_visibility(filename_t, filename_ds):
"""
Prints and handles the Visibility coverage mapping functionality.
:param filename_t:
@ -310,7 +318,7 @@ def menu_visibility(filename_t, filename_ds):
:return:
"""
global filter_applicable_to
clear()
_clear()
print('Menu: %s' % MENU_NAME_VISIBILITY_MAPPING)
print('')
print('Selected techniques YAML file: %s' % filename_t)
@ -326,40 +334,40 @@ def menu_visibility(filename_t, filename_ds):
print('4. Generate an Excel sheet with all administrated techniques.')
print('5. Check the technique YAML file for errors.')
print('9. Back to main menu.')
choice = ask_input()
choice = _ask_input()
if choice == '1':
print('Specify your filter for the applicable_to field:')
filter_applicable_to = ask_input().lower()
filter_applicable_to = _ask_input().lower()
elif choice == '2':
print('Writing visibility coverage layer...')
generate_visibility_layer(filename_t, filename_ds, False, filter_applicable_to)
wait()
_wait()
elif choice == '3':
print('Writing visibility coverage layer overlaid with detections...')
generate_visibility_layer(filename_t, filename_ds, True, filter_applicable_to)
wait()
_wait()
elif choice == '4':
print('Generating Excel file...')
export_techniques_list_to_excel(filename_t)
wait()
_wait()
elif choice == '5':
print('Checking the technique YAML file for errors...')
check_yaml_file_health(filename_t, FILE_TYPE_TECHNIQUE_ADMINISTRATION, health_is_called=True)
wait()
_wait()
elif choice == '9':
interactive_menu()
elif choice == 'q':
quit()
menu_visibility(filename_t, filename_ds)
_menu_visibility(filename_t, filename_ds)
def menu_groups():
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
clear()
_clear()
print('Menu: %s' % MENU_NAME_THREAT_ACTOR_GROUP_MAPPING)
print('')
print('Options:')
@ -375,22 +383,22 @@ def menu_groups():
print('')
print('7. Generate a heat map layer.')
print('9. Back to main menu.')
choice = ask_input()
choice = _ask_input()
if choice == '1':
print('Specify True or False for software group:')
software_group = True if ask_input().lower() == 'true' else False
software_group = True if _ask_input().lower() == 'true' else False
elif choice == '2':
print('Specify platform (all, Linux, macOS, Windows):')
p = ask_input().lower()
p = _ask_input().lower()
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()
s = _ask_input().lower()
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)')
g = ask_input()
g = _ask_input()
groups = g if g is not '' else 'all'
elif choice == '5':
print('')
@ -398,30 +406,30 @@ def menu_groups():
print('2. Overlay with detections.')
print('3. Overlay with visibility.')
print('4. No overlay.')
choice = ask_input()
choice = _ask_input()
if choice == '1':
print('Specify the group(s) to overlay (in a different color) on the one specified in the Groups option. '
'A group can be their ID, name or alias separated using commas. Other option is to provide a YAML '
'file with a custom group(s).')
overlay_type = OVERLAY_TYPE_GROUP
groups_overlay = ask_input()
groups_overlay = _ask_input()
elif choice == '2':
overlay_type = OVERLAY_TYPE_DETECTION
groups_overlay = select_file(MENU_NAME_THREAT_ACTOR_GROUP_MAPPING, 'techniques', FILE_TYPE_TECHNIQUE_ADMINISTRATION, False)
groups_overlay = _select_file(MENU_NAME_THREAT_ACTOR_GROUP_MAPPING, 'techniques', FILE_TYPE_TECHNIQUE_ADMINISTRATION, False)
elif choice == '3':
overlay_type = OVERLAY_TYPE_VISIBILITY
groups_overlay = select_file(MENU_NAME_THREAT_ACTOR_GROUP_MAPPING, 'techniques', FILE_TYPE_TECHNIQUE_ADMINISTRATION, False)
groups_overlay = _select_file(MENU_NAME_THREAT_ACTOR_GROUP_MAPPING, 'techniques', FILE_TYPE_TECHNIQUE_ADMINISTRATION, False)
elif choice == '4':
overlay_type = ''
groups_overlay = ''
elif choice == '6':
print('Specify your filter for the applicable_to field:')
filter_applicable_to = ask_input().lower()
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)
wait()
_wait()
elif choice == '9':
interactive_menu()
elif choice == 'q':
quit()
menu_groups()
_menu_groups()