From 7dbcd594e4ee4b979ace6c7bb32d9b902b724eb0 Mon Sep 17 00:00:00 2001
From: Yevhen Zavhorodnii <yevhen.zavhorodnii@ticketmaster.co.uk>
Date: Mon, 10 Jun 2024 17:13:39 +0100
Subject: [PATCH] Proposal on using custom diagrams

---
 internal/threagile/config.go | 18 ++++++++++++++++++
 pkg/report/generate.go       |  6 ++++--
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/internal/threagile/config.go b/internal/threagile/config.go
index 451e3f97..b3587299 100644
--- a/internal/threagile/config.go
+++ b/internal/threagile/config.go
@@ -25,6 +25,8 @@ type Config struct {
 	KeyFolder    string
 
 	InputFile                   string
+	UseCustomDataFlowDiagram    bool
+	UseCustomDataAssetDiagram   bool
 	DataFlowDiagramFilenamePNG  string
 	DataAssetDiagramFilenamePNG string
 	DataFlowDiagramFilenameDOT  string
@@ -77,6 +79,8 @@ func (c *Config) Defaults(buildTimestamp string) *Config {
 		KeyFolder:    KeyDir,
 
 		InputFile:                   InputFile,
+		UseCustomDataFlowDiagram:    false,
+		UseCustomDataAssetDiagram:   false,
 		DataFlowDiagramFilenamePNG:  DataFlowDiagramFilenamePNG,
 		DataAssetDiagramFilenamePNG: DataAssetDiagramFilenamePNG,
 		DataFlowDiagramFilenameDOT:  DataFlowDiagramFilenameDOT,
@@ -241,6 +245,12 @@ func (c *Config) Merge(config Config, values map[string]any) {
 		case strings.ToLower("InputFile"):
 			c.InputFile = config.InputFile
 
+		case strings.ToLower("UseCustomDataAssetDiagram"):
+			c.UseCustomDataAssetDiagram = config.UseCustomDataAssetDiagram
+
+		case strings.ToLower("UseCustomDataFlowDiagram"):
+			c.UseCustomDataFlowDiagram = config.UseCustomDataFlowDiagram
+
 		case strings.ToLower("DataFlowDiagramFilenamePNG"):
 			c.DataFlowDiagramFilenamePNG = config.DataFlowDiagramFilenamePNG
 
@@ -425,6 +435,14 @@ func (c *Config) GetInputFile() string {
 	return c.InputFile
 }
 
+func (c *Config) GetUseCustomDataFlowDiagram() bool {
+	return c.UseCustomDataFlowDiagram
+}
+
+func (c *Config) GetUseCustomDataAssetDiagram() bool {
+	return c.UseCustomDataAssetDiagram
+}
+
 func (c *Config) GetDataFlowDiagramFilenamePNG() string {
 	return c.DataFlowDiagramFilenamePNG
 }
diff --git a/pkg/report/generate.go b/pkg/report/generate.go
index 10bb2fc7..d7e85370 100644
--- a/pkg/report/generate.go
+++ b/pkg/report/generate.go
@@ -46,6 +46,8 @@ type reportConfigReader interface {
 	GetTempFolder() string
 
 	GetInputFile() string
+	GetUseCustomDataAssetDiagram() bool
+	GetUseCustomDataFlowDiagram() bool
 	GetDataFlowDiagramFilenamePNG() string
 	GetDataAssetDiagramFilenamePNG() string
 	GetDataFlowDiagramFilenameDOT() string
@@ -86,7 +88,7 @@ func Generate(config reportConfigReader, readResult *model.ReadResult, commands
 		diagramDPI = config.GetMaxGraphvizDPI()
 	}
 	// Data-flow Diagram rendering
-	if generateDataFlowDiagram {
+	if generateDataFlowDiagram && !config.GetUseCustomDataFlowDiagram() {
 		gvFile := filepath.Join(config.GetOutputFolder(), config.GetDataFlowDiagramFilenameDOT())
 		if !config.GetKeepDiagramSourceFiles() {
 			tmpFileGV, err := os.CreateTemp(config.GetTempFolder(), config.GetDataFlowDiagramFilenameDOT())
@@ -108,7 +110,7 @@ func Generate(config reportConfigReader, readResult *model.ReadResult, commands
 		}
 	}
 	// Data Asset Diagram rendering
-	if generateDataAssetsDiagram {
+	if generateDataAssetsDiagram && !config.GetUseCustomDataAssetDiagram() {
 		gvFile := filepath.Join(config.GetOutputFolder(), config.GetDataAssetDiagramFilenameDOT())
 		if !config.GetKeepDiagramSourceFiles() {
 			tmpFile, err := os.CreateTemp(config.GetTempFolder(), config.GetDataAssetDiagramFilenameDOT())