From 812b887c25af897aea1fa255b2d89d3161347be6 Mon Sep 17 00:00:00 2001
From: egibs <20933572+egibs@users.noreply.github.com>
Date: Fri, 20 Mar 2026 16:47:55 -0500
Subject: [PATCH] fix: improve JSON/YAML diff output; fix upgradeRisk edge case

Signed-off-by: egibs <20933572+egibs@users.noreply.github.com>
---
 pkg/render/json.go        | 16 +++++++++-------
 pkg/render/json_test.go   | 31 +++++++++++++++++++++++++++++++
 pkg/render/yaml.go        | 16 +++++++++-------
 pkg/render/yaml_test.go   | 31 +++++++++++++++++++++++++++++++
 pkg/report/report.go      | 28 +++++++++-------------------
 pkg/report/report_test.go |  2 ++
 6 files changed, 91 insertions(+), 33 deletions(-)

diff --git a/pkg/render/json.go b/pkg/render/json.go
index 5eeddb64a..f54000cd3 100644
--- a/pkg/render/json.go
+++ b/pkg/render/json.go
@@ -45,13 +45,15 @@ func (r JSON) Full(ctx context.Context, c *malcontent.Config, rep *malcontent.Re
 		Filter: "",
 	}
 
-	rep.Files.Range(func(key string, fr *malcontent.FileReport) bool {
-		if ctx.Err() != nil {
-			return false
-		}
-		sanitizeFileReport(key, fr, jr.Files)
-		return true
-	})
+	if rep.Files != nil {
+		rep.Files.Range(func(key string, fr *malcontent.FileReport) bool {
+			if ctx.Err() != nil {
+				return false
+			}
+			sanitizeFileReport(key, fr, jr.Files)
+			return true
+		})
+	}
 
 	if c != nil && c.Stats && jr.Diff == nil {
 		if s := serializedStats(c, rep); s != nil {
diff --git a/pkg/render/json_test.go b/pkg/render/json_test.go
index 38652f749..5aa139a49 100644
--- a/pkg/render/json_test.go
+++ b/pkg/render/json_test.go
@@ -256,6 +256,37 @@ func TestJSONRendererScanningNoOp(t *testing.T) {
 	}
 }
 
+func TestJSONRendererDiffWithNilFiles(t *testing.T) {
+	t.Parallel()
+	var buf bytes.Buffer
+	renderer := NewJSON(&buf)
+
+	ctx := context.Background()
+	cfg := &malcontent.Config{}
+	diff := &malcontent.DiffReport{
+		Added:    orderedmap.New[string, *malcontent.FileReport](),
+		Removed:  orderedmap.New[string, *malcontent.FileReport](),
+		Modified: orderedmap.New[string, *malcontent.FileReport](),
+	}
+	diff.Added.Set("/bin/added", &malcontent.FileReport{Path: "/bin/added", RiskScore: 2})
+	// Mirror what action.Diff() actually returns: Diff set, Files nil
+	report := &malcontent.Report{Diff: diff}
+
+	err := renderer.Full(ctx, cfg, report)
+	if err != nil {
+		t.Fatalf("Full() error = %v", err)
+	}
+
+	var result Report
+	if err := json.Unmarshal(buf.Bytes(), &result); err != nil {
+		t.Fatalf("Generated invalid JSON: %v", err)
+	}
+
+	if result.Diff == nil {
+		t.Error("Expected diff in output")
+	}
+}
+
 func TestJSONRendererFileNoOp(t *testing.T) {
 	t.Parallel()
 	var buf bytes.Buffer
diff --git a/pkg/render/yaml.go b/pkg/render/yaml.go
index 29232e9e9..f2df1e508 100644
--- a/pkg/render/yaml.go
+++ b/pkg/render/yaml.go
@@ -45,13 +45,15 @@ func (r YAML) Full(ctx context.Context, c *malcontent.Config, rep *malcontent.Re
 		Filter: "",
 	}
 
-	rep.Files.Range(func(key string, fr *malcontent.FileReport) bool {
-		if ctx.Err() != nil {
-			return false
-		}
-		sanitizeFileReport(key, fr, yr.Files)
-		return true
-	})
+	if rep.Files != nil {
+		rep.Files.Range(func(key string, fr *malcontent.FileReport) bool {
+			if ctx.Err() != nil {
+				return false
+			}
+			sanitizeFileReport(key, fr, yr.Files)
+			return true
+		})
+	}
 
 	if c != nil && c.Stats && yr.Diff == nil {
 		if s := serializedStats(c, rep); s != nil {
diff --git a/pkg/render/yaml_test.go b/pkg/render/yaml_test.go
index 4c7c2cad2..b7b68284f 100644
--- a/pkg/render/yaml_test.go
+++ b/pkg/render/yaml_test.go
@@ -256,6 +256,37 @@ func TestYAMLRendererScanningNoOp(t *testing.T) {
 	}
 }
 
+func TestYAMLRendererDiffWithNilFiles(t *testing.T) {
+	t.Parallel()
+	var buf bytes.Buffer
+	renderer := NewYAML(&buf)
+
+	ctx := context.Background()
+	cfg := &malcontent.Config{}
+	diff := &malcontent.DiffReport{
+		Added:    orderedmap.New[string, *malcontent.FileReport](),
+		Removed:  orderedmap.New[string, *malcontent.FileReport](),
+		Modified: orderedmap.New[string, *malcontent.FileReport](),
+	}
+	diff.Added.Set("/bin/added", &malcontent.FileReport{Path: "/bin/added", RiskScore: 2})
+	// Mirror what action.Diff() actually returns: Diff set, Files nil
+	report := &malcontent.Report{Diff: diff}
+
+	err := renderer.Full(ctx, cfg, report)
+	if err != nil {
+		t.Fatalf("Full() error = %v", err)
+	}
+
+	var result Report
+	if err := yaml.Unmarshal(buf.Bytes(), &result); err != nil {
+		t.Fatalf("Generated invalid YAML: %v", err)
+	}
+
+	if result.Diff == nil {
+		t.Error("Expected diff in output")
+	}
+}
+
 func TestYAMLRendererFileNoOp(t *testing.T) {
 	t.Parallel()
 	var buf bytes.Buffer
diff --git a/pkg/report/report.go b/pkg/report/report.go
index c4e156391..63d32fd89 100644
--- a/pkg/report/report.go
+++ b/pkg/report/report.go
@@ -836,26 +836,16 @@ func upgradeRisk(ctx context.Context, riskScore int, riskCounts map[int]int, siz
 	}
 	highCount := riskCounts[HIGH]
 	sizeMB := size / 1024 / 1024
-	upgrade := false
 
-	switch {
-	// small scripts, tiny ELF binaries
-	case size < 1024 && highCount > 1:
-		upgrade = true
-	// include most UPX binaries
-	case sizeMB < 2 && highCount > 2:
-		upgrade = true
-	case sizeMB < 4 && highCount > 3:
-		upgrade = true
-	case sizeMB < 10 && highCount > 4:
-		upgrade = true
-	case sizeMB < 20 && highCount > 5:
-		upgrade = true
-	case highCount > 6:
-		upgrade = true
-	}
-
-	clog.DebugContextf(ctx, "upgrading risk: high=%d, size=%d", highCount, size)
+	upgrade := (size < 1024 && highCount > 1) || // small scripts, tiny ELF binaries
+		(sizeMB < 2 && highCount > 2) || // include most UPX binaries
+		(sizeMB < 4 && highCount > 3) ||
+		(sizeMB < 10 && highCount > 4) ||
+		highCount > 5
+
+	if upgrade {
+		clog.DebugContextf(ctx, "upgrading risk to critical: high=%d, size=%d", highCount, size)
+	}
 	return upgrade
 }
 
diff --git a/pkg/report/report_test.go b/pkg/report/report_test.go
index 67ac65396..20e68e9bc 100644
--- a/pkg/report/report_test.go
+++ b/pkg/report/report_test.go
@@ -85,6 +85,8 @@ func TestUpgradeRisk(t *testing.T) {
 		{"small-risky", 3, map[int]int{3: 3}, 8192, true},
 		{"large-not", 3, map[int]int{3: 3}, 1024 * 1024 * 1024, false},
 		{"large-yes", 3, map[int]int{3: 10}, 1024 * 1024 * 1024, true},
+		{"large-default-threshold", 3, map[int]int{3: 6}, 1024 * 1024 * 1024, true},
+		{"large-below-threshold", 3, map[int]int{3: 5}, 1024 * 1024 * 1024, false},
 	}
 
 	for _, tt := range tests {