feat: Multiple DFDs per TM are now supported #29

threatcl · Mar 19, 2023 · 04c31db · 04c31db
1 parent af8957c
commit 04c31db
Show file tree

Hide file tree

Showing 13 changed files with 751 additions and 228 deletions.
diff --git a/cmd/hcltm/dashboard.go b/cmd/hcltm/dashboard.go
@@ -219,7 +219,10 @@ func (c *DashboardCommand) Run(args []string) int {
 			outfiles = append(outfiles, outfile)
 
 			if !c.flagNoDfd && len(tm.DataFlowDiagrams) > 0 {
-				outfiles = append(outfiles, outfilePath(c.flagOutDir, tm.Name, file, ".png"))
+				for _, adfd := range tm.DataFlowDiagrams {
+					dfdFile := outfilePath(c.flagOutDir, fmt.Sprintf("%s_%s", tm.Name, adfd.Name), file, ".png")
+					outfiles = append(outfiles, dfdFile)
+				}
 			}
 
 		}
@@ -255,24 +258,20 @@ func (c *DashboardCommand) Run(args []string) int {
 
 			// First we check if there are any DFDs
 			if !c.flagNoDfd && len(tm.DataFlowDiagrams) > 0 {
-				dfdPath := outfilePath(c.flagOutDir, tm.Name, file, ".png")
-				err = tm.GenerateDfdPng(dfdPath)
-				if err != nil {
-					fmt.Printf("Error generating DFD: %s\n", err)
-					return 1
-				}
+				for _, adfd := range tm.DataFlowDiagrams {
+					dfdPath := outfilePath(c.flagOutDir, fmt.Sprintf("%s_%s", tm.Name, adfd.Name), file, ".png")
+					err = adfd.GenerateDfdPng(dfdPath, tm.Name)
+					if err != nil {
+						fmt.Printf("Error generating DFD: %s\n", err)
+						return 1
+					}
 
-				fmt.Printf("Successfully wrote to '%s'\n", dfdPath)
+					fmt.Printf("Successfully wrote to '%s'\n", dfdPath)
 
-				// Now we set the tm diagram to this file
-				// if it's currently unset
-				if tm.DiagramLink == "" {
-					tm.DiagramLink = filepath.Base(dfdPath)
+					tm.AllDiagrams = append(tm.AllDiagrams, filepath.Base(dfdPath))
 				}
-
 			}
 
-			// tmBuffer, err := tm.RenderMarkdown(spec.TmMDTemplate)
 			tmBuffer, err := tm.RenderMarkdown(tmTemplate)
 			if err != nil {
 				fmt.Println(err)

diff --git a/cmd/hcltm/dashboard_test.go b/cmd/hcltm/dashboard_test.go
@@ -154,7 +154,7 @@ func TestDashboardWithDfd(t *testing.T) {
 		t.Errorf("Expected %s to contain %s", dbfile, "tm3-tm2one.md")
 	}
 
-	f, err := os.Open(fmt.Sprintf("%s/out/tm3-tm2one.png", d))
+	f, err := os.Open(fmt.Sprintf("%s/out/tm3-tm2onelegacydfd.png", d))
 	if err != nil {
 		t.Fatalf("Error opening png: %s", err)
 	}