Add --force-overwrite flag to force overwriting the file we specify in --output-filename (#124)

* Add --force-overwrite flag to force overwriting the file we specify in --output-filename

* Add use_existing_filename() to improve readability and handle duplicate extension issues

* Make use_existing_filename() code more consistent with get_non_existing_filename()

* local stix path for refresh script

* fixes for PR #123

---------

Co-authored-by: Ruben Bouman <ruben@siriussecurity.nl>
Co-authored-by: Ruben Bouman <11735227+rubinatorz@users.noreply.github.com>

Author: alexhaydock

data_source_mapping.py

@@ -196,11 +196,12 @@ def _get_technique_yaml_obj(techniques, tech_id):
             return tech
 
 
-def generate_data_sources_layer(filename, output_filename, layer_name, layer_settings):
+def generate_data_sources_layer(filename, output_filename, output_overwrite, layer_name, layer_settings):
     """
     Generates a generic layer for data sources.
     :param filename: the filename of the YAML file containing the data sources administration
     :param output_filename: the output filename defined by the user
+    :param output_overwrite: boolean flag indicating whether we're in overwrite mode
     :param layer_name: the name of the Navigator layer
     :param layer_settings: settings for the Navigator layer
     :return:
@@ -220,14 +221,15 @@ def generate_data_sources_layer(filename, output_filename, layer_name, layer_set
     json_string = simplejson.dumps(layer).replace('}, ', '},\n')
     if not output_filename:
         output_filename = create_output_filename('data_sources', name)
-    write_file(output_filename, json_string)
+    write_file(output_filename, output_overwrite, json_string)
 
 
-def plot_data_sources_graph(filename, output_filename):
+def plot_data_sources_graph(filename, output_filename, output_overwrite):
     """
     Generates a line graph which shows the improvements on numbers of data sources through time.
     :param filename: the filename of the YAML file containing the data sources administration
     :param output_filename: the output filename defined by the user
+    :param output_overwrite: boolean flag indicating whether we're in overwrite mode
     :return:
     """
     my_data_sources, name, _, _, _ = load_data_sources(filename)
@@ -247,7 +249,11 @@ def plot_data_sources_graph(filename, output_filename):
         output_filename = 'graph_data_sources'
     elif output_filename.endswith('.html'):
         output_filename = output_filename.replace('.html', '')
-    output_filename = get_non_existing_filename('output/' + output_filename, 'html')
+
+    if not output_overwrite:
+        output_filename = get_non_existing_filename('output/' + output_filename, 'html')
+    else:
+        output_filename = use_existing_filename('output/' + output_filename, 'html')
 
     import plotly.graph_objs as go
     import plotly.offline as offline
@@ -259,23 +265,30 @@ def plot_data_sources_graph(filename, output_filename):
     print("File written:   " + output_filename)
 
 
-def export_data_source_list_to_excel(filename, output_filename, eql_search=False):
+def export_data_source_list_to_excel(filename, output_filename, output_overwrite, eql_search=False):
     """
     Makes an overview of all MITRE ATT&CK data sources (via techniques) and lists which data sources are present
     in the YAML administration including all properties and data quality score.
     :param filename: the filename of the YAML file containing the data sources administration
     :param output_filename: the output filename defined by the user
+    :param output_overwrite: boolean flag indicating whether we're in overwrite mode
     :param eql_search: specify if an EQL search was performed which may have resulted in missing ATT&CK data sources
     :return:
     """
     # pylint: disable=unused-variable
     my_data_sources, name, systems, _, domain = load_data_sources(filename, filter_empty_scores=False)
     my_data_sources = dict(sorted(my_data_sources.items(), key=lambda kv: kv[0], reverse=False))
+
     if not output_filename:
         output_filename = 'data_sources'
     elif output_filename.endswith('.xlsx'):
         output_filename = output_filename.replace('.xlsx', '')
-    excel_filename = get_non_existing_filename('output/' + output_filename, 'xlsx')
+
+    if not output_overwrite:
+        excel_filename = get_non_existing_filename('output/' + output_filename, 'xlsx')
+    else:
+        excel_filename = use_existing_filename('output/' + output_filename, 'xlsx')
+
     workbook = xlsxwriter.Workbook(excel_filename)
     worksheet = workbook.add_worksheet('Data sources')
 
@@ -868,11 +881,12 @@ def update_technique_administration_file(file_data_sources, file_tech_admin):
 # pylint: disable=redefined-outer-name
 
 
-def generate_technique_administration_file(filename, output_filename, write_file=True, all_techniques=False):
+def generate_technique_administration_file(filename, output_filename, output_overwrite, write_file=True, all_techniques=False):
     """
     Generate a technique administration file based on the data source administration YAML file
     :param filename: the filename of the YAML file containing the data sources administration
     :param output_filename: the output filename defined by the user
+    :param output_overwrite: boolean flag indicating whether we're in overwrite mode
     :param write_file: by default the file is written to disk
     :param all_techniques: include all ATT&CK techniques in the generated YAML file that are applicable to the
     platform(s) specified in the data source YAML file
@@ -985,7 +999,12 @@ def generate_technique_administration_file(filename, output_filename, write_file
             output_filename = 'techniques-administration-' + normalize_name_to_filename(name)
         elif output_filename.endswith('.yaml'):
             output_filename = output_filename.replace('.yaml', '')
-        output_filename = get_non_existing_filename('output/' + output_filename, 'yaml')
+
+        if not output_overwrite:
+            output_filename = get_non_existing_filename(f'output/{output_filename}', 'yaml')
+        else:
+            output_filename = use_existing_filename(f'output/{output_filename}', 'yaml')
+
         with open(output_filename, 'w') as f:
             f.writelines(yaml_file_lines)
         print("File written:   " + output_filename)

dettect.py

@@ -70,6 +70,8 @@ def _init_menu():
                                                             'scores are calculated in the same way as with the option: '
                                                             '-y, --yaml', action='store_true')
     parser_data_sources.add_argument('-of', '--output-filename', help='set the output filename')
+    parser_data_sources.add_argument('--force-overwrite', help='force overwriting the output file if it already exists',
+                                     action='store_true')
     parser_data_sources.add_argument('-ln', '--layer-name', help='set the name of the Navigator layer')
     parser_data_sources.add_argument('--health', help='check the YAML file(s) for errors', action='store_true')
     parser_data_sources.add_argument('--local-stix-path', help='path to a local STIX repository to use DeTT&CT offline '
@@ -109,6 +111,8 @@ def _init_menu():
     parser_visibility.add_argument('-g', '--graph', help='generate a graph with visibility added through time',
                                    action='store_true')
     parser_visibility.add_argument('-of', '--output-filename', help='set the output filename')
+    parser_visibility.add_argument('--force-overwrite', help='force overwriting the output file if it already exists',
+                                     action='store_true')
     parser_visibility.add_argument('-ln', '--layer-name', help='set the name of the Navigator layer')
     parser_visibility.add_argument('-cd', '--count-detections', help='Show the number of detections instead of listing '
                                    'all detection locations in Layer metadata (when using '
@@ -155,6 +159,8 @@ def _init_menu():
     parser_detection.add_argument('-g', '--graph', help='generate a graph with detections added through time',
                                   action='store_true')
     parser_detection.add_argument('-of', '--output-filename', help='set the output filename')
+    parser_detection.add_argument('--force-overwrite', help='force overwriting the output file if it already exists',
+                                     action='store_true')
     parser_detection.add_argument('-ln', '--layer-name', help='set the name of the Navigator layer')
     parser_detection.add_argument('-cd', '--count-detections', help='Show the number of detections instead of listing '
                                                                     'all detection locations in Layer metadata. Location '
@@ -224,6 +230,8 @@ def _init_menu():
                                                    'most recent \'score\' objects',
                               action='store_true', default=False)
     parser_group.add_argument('-of', '--output-filename', help='set the output filename')
+    parser_group.add_argument('--force-overwrite', help='force overwriting the output file if it already exists',
+                                     action='store_true')
     parser_group.add_argument('-ln', '--layer-name', help='set the name of the Navigator layer')
     parser_group.add_argument('-cd', '--count-detections', help='Show the number of detections instead of listing '
                               'all detection locations in Layer metadata (when using an overlay with detection). Location '
@@ -301,13 +309,13 @@ def _menu(menu_parser):
             if args.update and check_file(args.file_tech, FILE_TYPE_TECHNIQUE_ADMINISTRATION, args.health):
                 update_technique_administration_file(file_ds, args.file_tech)
             if args.layer:
-                generate_data_sources_layer(file_ds, args.output_filename, args.layer_name, layer_settings)
+                generate_data_sources_layer(file_ds, args.output_filename, args.force_overwrite, args.layer_name, layer_settings)
             if args.excel:
-                export_data_source_list_to_excel(file_ds, args.output_filename, eql_search=args.search)
+                export_data_source_list_to_excel(file_ds, args.output_filename, args.force_overwrite, eql_search=args.search)
             if args.graph:
-                plot_data_sources_graph(file_ds, args.output_filename)
+                plot_data_sources_graph(file_ds, args.output_filename, args.force_overwrite)
             if args.yaml:
-                generate_technique_administration_file(file_ds, args.output_filename, all_techniques=args.yaml_all_techniques)
+                generate_technique_administration_file(file_ds, args.output_filename, args.force_overwrite, all_techniques=args.yaml_all_techniques)
 
     elif args.subparser in ['visibility', 'v']:
         if check_file(args.file_tech, FILE_TYPE_TECHNIQUE_ADMINISTRATION, args.health):
@@ -323,22 +331,22 @@ def _menu(menu_parser):
                 if not file_tech:
                     quit()  # something went wrong in executing the search or 0 results where returned
             if args.layer:
-                generate_visibility_layer(file_tech, False, args.output_filename, args.layer_name,
+                generate_visibility_layer(file_tech, False, args.output_filename, args.force_overwrite, args.layer_name,
                                           layer_settings, args.platform, args.count_detections)
             if args.overlay:
-                generate_visibility_layer(file_tech, True, args.output_filename, args.layer_name,
+                generate_visibility_layer(file_tech, True, args.output_filename, args.force_overwrite, args.layer_name,
                                           layer_settings, args.platform, args.count_detections)
             if args.graph:
-                plot_graph(file_tech, 'visibility', args.output_filename)
+                plot_graph(file_tech, 'visibility', args.output_filename, args.force_overwrite)
             if args.excel:
-                export_techniques_list_to_excel(file_tech, args.output_filename)
+                export_techniques_list_to_excel(file_tech, args.output_filename, args.force_overwrite)
 
     # TODO add Group EQL search capabilities
     elif args.subparser in ['group', 'g']:
         layer_settings = _parse_layer_settings(args.layer_settings)
         generate_group_heat_map(args.groups, args.campaigns, args.overlay, args.overlay_type, args.platform,
                                 args.software, args.include_software, args.search_visibility, args.search_detection, args.health,
-                                args.output_filename, args.layer_name, args.domain, layer_settings,
+                                args.output_filename, args.force_overwrite, args.layer_name, args.domain, layer_settings,
                                 args.all_scores, args.count_detections)
 
     elif args.subparser in ['detection', 'd']:
@@ -355,14 +363,15 @@ def _menu(menu_parser):
                 if not file_tech:
                     quit()  # something went wrong in executing the search or 0 results where returned
             if args.layer:
-                generate_detection_layer(file_tech, False, args.output_filename, args.layer_name,
+                generate_detection_layer(file_tech, False, args.output_filename, args.force_overwrite, args.layer_name,
                                          layer_settings, args.platform, args.count_detections)
             if args.overlay:
-                generate_detection_layer(file_tech, True, args.output_filename, args.layer_name, layer_settings, args.platform, args.count_detections)
+                generate_detection_layer(file_tech, True, args.output_filename, args.force_overwrite, args.layer_name,
+                                         layer_settings, args.platform, args.count_detections)
             if args.graph:
-                plot_graph(file_tech, 'detection', args.output_filename)
+                plot_graph(file_tech, 'detection', args.output_filename, args.force_overwrite)
             if args.excel:
-                export_techniques_list_to_excel(file_tech, args.output_filename)
+                export_techniques_list_to_excel(file_tech, args.output_filename, args.force_overwrite)
 
     elif args.subparser in ['generic', 'ge']:
         if args.datasources:

file_output.py

@@ -11,16 +11,21 @@ def _clean_filename(filename):
     return filename.replace('/', '').replace('\\', '').replace(':', '')[:200]
 
 
-def write_file(filename, content):
+def write_file(filename, overwrite_mode, content):
     """
     Writes content to a file and ensures if the file already exists it won't be overwritten by appending a number
     as suffix.
     :param filename: filename
+    :param overwrite_mode: defines whether we want to force overwriting existing file
     :param content: the content of the file that needs to be written to the file
     :return:
     """
     output_filename = 'output/%s' % _clean_filename(filename)
-    output_filename = get_non_existing_filename(output_filename, 'json')
+
+    if not overwrite_mode:
+        output_filename = get_non_existing_filename(output_filename, 'json')
+    else:
+        output_filename = use_existing_filename(output_filename, 'json')
 
     with open(output_filename, 'w') as f:
         f.write(content)
@@ -57,9 +62,9 @@ def create_output_filename(filename_prefix, filename):
 def get_non_existing_filename(filename, extension):
     """
     Generates a filename that doesn't exist based on the given filename by appending a number as suffix.
-    :param filename:
-    :param extension:
-    :return:
+    :param filename: input filename
+    :param extension: input extension
+    :return: unique filename
     """
     if filename.endswith('.' + extension):
         filename = filename.replace('.' + extension, '')
@@ -73,6 +78,20 @@ def get_non_existing_filename(filename, extension):
     return output_filename
 
 
+def use_existing_filename(filename, extension):
+    """
+    Generates a filename that preserves the file extension if present.
+    If no extension is present, adds the provided extension.
+    :param filename: input filename
+    :param extension: input extension
+    :return: filename and extension, without duplicating extensions
+    """
+    if filename.endswith('.' + extension):
+        filename = filename.replace('.' + extension, '')
+    output_filename = '%s.%s' % (filename, extension)
+    return output_filename
+
+
 def normalize_name_to_filename(name):
     """
     Normalize the input filename to a lowercase filename and replace spaces with dashes.

group_mapping.py

@@ -604,8 +604,8 @@ def _get_group_list(groups, file_type):
 
 
 def generate_group_heat_map(groups, campaigns, overlay, overlay_type, platform, overlay_software, include_software,
-                            search_visibility, search_detection, health_is_called, output_filename, layer_name, domain,
-                            layer_settings, include_all_score_objs, count_detections):
+                            search_visibility, search_detection, health_is_called, output_filename, output_overwrite,
+                            layer_name, domain, layer_settings, include_all_score_objs, count_detections):
     """
     Calls all functions that are necessary for the generation of the heat map and write a json layer to disk.
     :param groups: threat actor groups
@@ -620,6 +620,7 @@ def generate_group_heat_map(groups, campaigns, overlay, overlay_type, platform,
     :param search_detection: detection EQL search query
     :param health_is_called: boolean that specifies if detailed errors in the file will be printed
     :param output_filename: output filename defined by the user
+    :param output_overwrite: boolean flag indicating whether we're in overwrite mode
     :param layer_name: the name of the Navigator layer
     :param domain: the specified domain
     :param layer_settings: settings for the Navigator layer
@@ -800,6 +801,6 @@ def generate_group_heat_map(groups, campaigns, overlay, overlay_type, platform,
             filename += '-overlay_' + '_'.join(overlay_list)
 
         filename = create_output_filename('attack', filename)
-        write_file(filename, json_string)
+        write_file(filename, output_overwrite, json_string)
     else:
-        write_file(output_filename, json_string)
+        write_file(output_filename, output_overwrite, json_string)