Skip to content

Commit

Permalink
test: updated tests for DOT output
Browse files Browse the repository at this point in the history
  • Loading branch information
@xntrik @teamcfr
xntrik authored and teamcfr committed Feb 6, 2023
1 parent f799e73 commit 8e0e624
Show file tree
Hide file tree
Showing 3 changed files with 113 additions and 188 deletions.
45 changes: 45 additions & 0 deletions cmd/hcltm/dfd_test.go
View file
Expand Up @@ -79,6 +79,51 @@ func TestDfd(t *testing.T) {

}

func TestDfdSvg(t *testing.T) {
d, err := ioutil.TempDir("", "")
if err != nil {
t.Fatalf("Error creatig tmp dir: %s", err)
}

defer os.RemoveAll(d)

cmd := testDfdCommand(t)

var code int

out := capturer.CaptureStdout(func() {
code = cmd.Run([]string{
fmt.Sprintf("-outdir=%s/out", d),
"-svg",
"./testdata/tm3.hcl",
})
})

if code != 0 {
t.Errorf("Code did not equal 0: %d", code)
}

if !strings.Contains(out, "Successfully created") {
t.Errorf("%s did not contain %s", out, "Successfully created")
}

f, err := os.Open(fmt.Sprintf("%s/out/tm3-tm2one.svg", d))
if err != nil {
t.Fatalf("Error opening svg: %s", err)
}

buffer := make([]byte, 512)
_, err = f.Read(buffer)
if err != nil {
t.Fatalf("Error reading svg: %s", err)
}

if http.DetectContentType(buffer) != "text/xml; charset=utf-8" {
t.Errorf("The output file isn't a svg, it's a '%s'", http.DetectContentType(buffer))
}

}

func TestDfdDot(t *testing.T) {
d, err := ioutil.TempDir("", "")
if err != nil {
Expand Down
199 changes: 15 additions & 184 deletions pkg/spec/dfd.go
View file
Expand Up @@ -11,20 +11,21 @@ import (
"gonum.org/v1/gonum/graph/encoding"
)

func (tm *Threatmodel) GenerateDot() (string, error) {
tmpFile, err := ioutil.TempFile("", "dot")
if err != nil {
return "", err
}
defer os.RemoveAll(tmpFile.Name())

// dot, err := tm.generateDfdDotFile(tmpFile.Name())
dot, err := tm.DataFlowDiagrams[0].generateDfdDotFile(tmpFile.Name(), tm.Name)
if err != nil {
return "", err
}
return dot, nil
}
// Commenting the below out - I don't think it's used anymore
// func (tm *Threatmodel) GenerateDot() (string, error) {
// tmpFile, err := ioutil.TempFile("", "dot")
// if err != nil {
// return "", err
// }
// defer os.RemoveAll(tmpFile.Name())
//
// // dot, err := tm.generateDfdDotFile(tmpFile.Name())
// dot, err := tm.DataFlowDiagrams[0].generateDfdDotFile(tmpFile.Name(), tm.Name)
// if err != nil {
// return "", err
// }
// return dot, nil
// }

func (d *DataFlowDiagram) GenerateDot(tmName string) (string, error) {
tmpFile, err := ioutil.TempFile("", "dot")
Expand Down Expand Up @@ -310,176 +311,6 @@ func (d *DataFlowDiagram) generateDfdDotFile(filepath, tmName string) (string, e
return dot, nil
}

// func (tm *Threatmodel) generateDfdDotFile(filepath string) (string, error) {
// // Build the DFD
// g := dfd.InitializeDFD(tm.Name)
//
// zones := make(map[string]*dfd.TrustBoundary)
// processes := make(map[string]*dfd.Process)
// external_elements := make(map[string]*dfd.ExternalService)
// data_stores := make(map[string]*dfd.DataStore)
//
// // Add zones
// for _, zone := range tm.DataFlowDiagram.TrustZones {
// if _, existing := zones[zone.Name]; !existing {
// newZone, err := g.AddTrustBoundary(zone.Name, "red")
// zones[zone.Name] = newZone
// if err != nil {
// return "", err
// }
// }
//
// // Add Processes from inside zone
// for _, process := range zone.Processes {
// err, newProcess := newDfdProcess(process.Name)
// if err != nil {
// return "", err
// }
// processes[process.Name] = newProcess
// zones[zone.Name].AddNodeElem(processes[process.Name])
// }
//
// // Add External Elements from inside zone
// for _, external_element := range zone.ExternalElements {
// err, newElement := newDfdExternalEntity(external_element.Name)
// if err != nil {
// return "", err
// }
// external_elements[external_element.Name] = newElement
// zones[zone.Name].AddNodeElem(external_elements[external_element.Name])
// }
//
// // Add Data Stores from inside zone
// for _, data_store := range zone.DataStores {
// err, newStore := newDfdStore(data_store.Name)
// if err != nil {
// return "", err
// }
// data_stores[data_store.Name] = newStore
// zones[zone.Name].AddNodeElem(data_stores[data_store.Name])
// }
//
// }
//
// // Add Processes
// for _, process := range tm.DataFlowDiagram.Processes {
// err, newProcess := newDfdProcess(process.Name)
// if err != nil {
// return "", err
// }
// processes[process.Name] = newProcess
//
// if process.TrustZone != "" {
// if _, ok := zones[process.TrustZone]; !ok {
// zone, err := g.AddTrustBoundary(process.TrustZone, "red")
// zones[process.TrustZone] = zone
// if err != nil {
// return "", err
// }
// }
//
// zones[process.TrustZone].AddNodeElem(processes[process.Name])
// } else {
// g.AddNodeElem(processes[process.Name])
// }
// }
//
// // Add External Elements
// for _, external_element := range tm.DataFlowDiagram.ExternalElements {
// err, newElement := newDfdExternalEntity(external_element.Name)
// if err != nil {
// return "", err
// }
// external_elements[external_element.Name] = newElement
//
// if external_element.TrustZone != "" {
// if _, ok := zones[external_element.TrustZone]; !ok {
// zone, err := g.AddTrustBoundary(external_element.TrustZone, "red")
// zones[external_element.TrustZone] = zone
// if err != nil {
// return "", err
// }
// }
//
// zones[external_element.TrustZone].AddNodeElem(external_elements[external_element.Name])
// } else {
// g.AddNodeElem(external_elements[external_element.Name])
// }
// }
//
// // Add Data Stores
// for _, data_store := range tm.DataFlowDiagram.DataStores {
// err, newStore := newDfdStore(data_store.Name)
// if err != nil {
// return "", err
// }
// data_stores[data_store.Name] = newStore
//
// if data_store.TrustZone != "" {
// if _, ok := zones[data_store.TrustZone]; !ok {
// zone, err := g.AddTrustBoundary(data_store.TrustZone, "red")
// zones[data_store.TrustZone] = zone
// if err != nil {
// return "", err
// }
// }
//
// zones[data_store.TrustZone].AddNodeElem(data_stores[data_store.Name])
// } else {
// g.AddNodeElem(data_stores[data_store.Name])
// }
// }
//
// for _, flow := range tm.DataFlowDiagram.Flows {
//
// var to, from graph.Node
//
// for name, process := range processes {
// if name == flow.From {
// from = process
// }
//
// if name == flow.To {
// to = process
// }
// }
//
// for name, external_element := range external_elements {
// if name == flow.From {
// from = external_element
// }
//
// if name == flow.To {
// to = external_element
// }
// }
//
// for name, data_store := range data_stores {
// if name == flow.From {
// from = data_store
// }
//
// if name == flow.To {
// to = data_store
// }
// }
//
// // g.AddFlow(processes[flow.From], processes[flow.To], flow.Name)
// g.AddFlow(from, to, flow.Name)
// }
//
// // Construct temp file for the dot file output
// // The library we use needs to save an actual file,
// // even though we don't use it, and instead use the raw
// // text output
// client := dfd.NewClient(filepath)
// dot, err := client.DFDToDOT(g)
// if err != nil {
// return "", err
// }
// return dot, nil
// }

func dotToPng(raw []byte, file string) error {
g, err := graphviz.ParseBytes(raw)
if err != nil {
Expand Down
57 changes: 53 additions & 4 deletions pkg/spec/dfd_test.go
View file
Expand Up @@ -243,11 +243,60 @@ func TestDfdPngGenerate(t *testing.T) {
}
}

func TestDfdSvgGenerate(t *testing.T) {
// tm := dfdTm()
//
// fulltm := fullDfdTm()
func TestDfdDotGenerate(t *testing.T) {
cases := []struct {
name string
tm *Threatmodel
exp string
errorthrown bool
}{
{
"valid_full_dfd",
fullDfdTm(),
"",
false,
},
{
"valid_full_dfd2",
fullDfdTm2(),
"",
false,
},
}

for _, tc := range cases {
tc := tc

t.Run(tc.name, func(t *testing.T) {
// t.Parallel()

d, err := ioutil.TempDir("", "")
if err != nil {
t.Fatalf("Error creating tmp dir: %s", err)
}
defer os.RemoveAll(d)

dot, err := tc.tm.DataFlowDiagrams[0].GenerateDot(tc.tm.Name)

if err != nil {
if !strings.Contains(err.Error(), tc.exp) {
t.Errorf("%s: Error rendering png: %s", tc.name, err)
}
} else {
if tc.errorthrown {
t.Errorf("%s: an error was thrown when it shoulnd't have", tc.name)
} else {
if !strings.Contains(dot, "graph") {
t.Errorf("%s: Could not find `graph` in DOT output", tc.name)
}
}
}

})
}
}

func TestDfdSvgGenerate(t *testing.T) {
cases := []struct {
name string
tm *Threatmodel
Expand Down

0 comments on commit 8e0e624

Please sign in to comment.