Added ATT&CK analyze options for colors and scores

sigma/cli/analyze.py

@@ -27,6 +27,31 @@ def analyze_group():
     "--subtechniques/--no-subtechniques", "-s/-S",
     default=True,
 )
+@click.option(
+    "--max-color", "-c",
+    default="#ff0000",
+    show_default=True,
+    help="Color used for maximum score."
+)
+@click.option(
+    "--min-color", "-C",
+    default="#ffffff00",
+    show_default=True,
+    help="Color used for zero score."
+)
+@click.option(
+    "--max-score", "-m",
+    type=int,
+    default=None,
+    help="Set fixed maximum score. All scores above are rendered as maximum. Increases color scale resolution for scores below.",
+)
+@click.option(
+    "--min-score", "-M",
+    type=int,
+    default="0",
+    show_default=True,
+    help="Minimum score. All scores below are not explicitly colored.",
+)
 @click.argument(
     "function",
     type=click.Choice(score_functions.keys()),
@@ -41,7 +66,7 @@ def analyze_group():
     required=True,
     type=click.Path(exists=True, allow_dash=True, path_type=pathlib.Path),
 )
-def analyze_attack(file_pattern, subtechniques, function, output, input):
+def analyze_attack(file_pattern, subtechniques, max_color, min_color, max_score, min_score, function, output, input):
     rules = load_rules(input, file_pattern)
     score_function = score_functions[function][0]
     scores = calculate_attack_scores(rules, score_function, not subtechniques)
@@ -71,11 +96,11 @@ def analyze_attack(file_pattern, subtechniques, function, output, input):
         "description": f"Sigma coverage heatmap generated by Sigma CLI with score function {function}",
         "gradient": {
             "colors": [
-                "#f5fff5",
-                "#00ff00",
+                min_color,
+                max_color,
             ],
-            "minValue": 0,
-            "maxValue": max(scores.values())
+            "minValue": min_score,
+            "maxValue": max_score or max(scores.values())
         },
         "techniques": layer_techniques,
     }

tests/test_analyze.py

@@ -25,6 +25,34 @@ def test_attack_generate():
     assert len(result.stdout.split()) > 25
     assert "T1505.003" in result.stdout
     assert "persistence" in result.stdout
+    assert "#ff0000" in result.stdout
+    assert "#ffffff00" in result.stdout
+    assert '"maxValue": 4' in result.stdout
+    assert '"minValue": 0' in result.stdout
+
+def test_attack_generate_max_value():
+    cli = CliRunner()
+    result = cli.invoke(analyze_attack, ["--max-score", "2", "max", "-", "tests/files/valid"])
+    assert result.exit_code == 0
+    assert '"maxValue": 2' in result.stdout
+
+def test_attack_generate_min_value():
+    cli = CliRunner()
+    result = cli.invoke(analyze_attack, ["--min-score", "2", "max", "-", "tests/files/valid"])
+    assert result.exit_code == 0
+    assert '"minValue": 2' in result.stdout
+
+def test_attack_generate_max_color():
+    cli = CliRunner()
+    result = cli.invoke(analyze_attack, ["--max-color", "#123456", "max", "-", "tests/files/valid"])
+    assert result.exit_code == 0
+    assert "#123456" in result.stdout
+
+def test_attack_generate_min_color():
+    cli = CliRunner()
+    result = cli.invoke(analyze_attack, ["--min-color", "#123456", "max", "-", "tests/files/valid"])
+    assert result.exit_code == 0
+    assert "#123456" in result.stdout
 
 def test_attack_generate_no_subtechniques():
     cli = CliRunner()