From 9b52bf1136e903b78a1d507edf01661f48984d5d Mon Sep 17 00:00:00 2001
From: Ruben Bouman <ruben@siriussecurity.nl>
Date: Thu, 2 May 2019 16:47:40 +0200
Subject: [PATCH] Improvend the legend for group + detection and visibility
 overlays.

---
 generic.py       | 28 ++++++++++++++++++----------
 group_mapping.py |  2 +-
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/generic.py b/generic.py
index 86bb5c9..4dc46e1 100644
--- a/generic.py
+++ b/generic.py
@@ -96,7 +96,7 @@ def _get_base_template(name, description, stage, platform, sorting):
     return layer
 
 
-def get_layer_template_groups(name, max_score, description, stage, platform):
+def get_layer_template_groups(name, max_score, description, stage, platform, overlay_type):
     """
     Prepares a base template for the json layer file that can be loaded into the MITRE ATT&CK Navigator.
     More information on the version 2.1 layer format:
@@ -106,19 +106,27 @@ def get_layer_template_groups(name, max_score, description, stage, platform):
     :param description: description
     :param stage: stage (act | prepare)
     :param platform: platform
+    :param overlay_type: group, visibility or detection
     :return: layer template dictionary
     """
     layer = _get_base_template(name, description, stage, platform, 3)
     layer['gradient'] = {'colors': [COLOR_GRADIENT_MIN, COLOR_GRADIENT_MAX], 'minValue': 0, 'maxValue': max_score}
-    layer['legendItems'] = \
-        [
-            {'label': 'Tech. ref. for ' + str(1) + ' group', 'color': COLOR_GRADIENT_MIN},
-            {'label': 'Tech. ref. for ' + str(max_score) + ' groups', 'color': COLOR_GRADIENT_MAX},
-            {'label': 'Groups overlay: tech. in group + overlay', 'color': COLOR_GROUP_OVERLAY_MATCH},
-            {'label': 'Groups overlay: tech. in overlay', 'color': COLOR_GROUP_OVERLAY_NO_MATCH},
-            {'label': 'Src. of tech. is only software', 'color': COLOR_SOFTWARE},
-            {'label': 'Src. of tech. is group(s)/overlay + software', 'color': COLOR_GROUP_AND_SOFTWARE}
-        ]
+    layer['legendItems'] = []
+    layer['legendItems'].append({'label': 'Tech. ref. for ' + str(1) + ' group', 'color': COLOR_GRADIENT_MIN})
+    layer['legendItems'].append({'label': 'Tech. ref. for ' + str(max_score) + ' groups', 'color': COLOR_GRADIENT_MAX})
+
+    if overlay_type == 'group':
+        layer['legendItems'].append({'label': 'Groups overlay: tech. in group + overlay', 'color': COLOR_GROUP_OVERLAY_MATCH})
+        layer['legendItems'].append({'label': 'Groups overlay: tech. in overlay', 'color': COLOR_GROUP_OVERLAY_NO_MATCH})
+        layer['legendItems'].append({'label': 'Src. of tech. is only software', 'color': COLOR_SOFTWARE})
+        layer['legendItems'].append({'label': 'Src. of tech. is group(s)/overlay + software', 'color': COLOR_GROUP_AND_SOFTWARE})
+    elif overlay_type == 'detection':
+        layer['legendItems'].append({'label': 'Tech. in group + detection', 'color': COLOR_GROUP_OVERLAY_MATCH})
+        layer['legendItems'].append({'label': 'Tech. in detection', 'color': COLOR_GROUP_OVERLAY_ONLY_DETECTION})
+    elif overlay_type == 'visibility':
+        layer['legendItems'].append({'label': 'Tech. in group + visibility', 'color': COLOR_GROUP_OVERLAY_MATCH})
+        layer['legendItems'].append({'label': 'Tech. in visibility', 'color': COLOR_GROUP_OVERLAY_ONLY_VISIBILITY})
+
     return layer
 
 
diff --git a/group_mapping.py b/group_mapping.py
index 6891294..d6f83c9 100644
--- a/group_mapping.py
+++ b/group_mapping.py
@@ -494,7 +494,7 @@ def generate_group_heat_map(groups, overlay, overlay_type, stage, platform, soft
     desc = 'stage: ' + stage + ' | platform: ' + platform + ' | group(s): ' + ', '.join(groups_list) + \
            ' | overlay group(s): ' + ', '.join(overlay_list)
 
-    layer = get_layer_template_groups(stage[0].upper() + stage[1:] + ' ' + platform, max_technique_count, desc, stage, platform)
+    layer = get_layer_template_groups(stage[0].upper() + stage[1:] + ' ' + platform, max_technique_count, desc, stage, platform, overlay_type)
     layer['techniques'] = technique_layer
 
     json_string = simplejson.dumps(layer).replace('}, ', '},\n')