Skip to content

Commit

Permalink
feat: Allow raw DOT output and start on better colour handling threat…
Browse files Browse the repository at this point in the history 
…cl#33 threatcl#27
  • Loading branch information
@xntrik
xntrik committed Jan 24, 2022
1 parent 4b38a9a commit d9f6eb2
Show file tree
Hide file tree
Showing 4 changed files with 184 additions and 52 deletions.
95 changes: 79 additions & 16 deletions cmd/hcltm/dfd.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,23 +15,30 @@ type DfdCommand struct {
flagOutDir string
flagOutFile string
flagOverwrite bool
flagDot bool
}

func (c *DfdCommand) Help() string {
helpText := `
Usage: hcltm dfd [options] -outdir=<directory> <files>
Generate Data Flow Diagram PNG files from existing Threat model HCL files
Generate Data Flow Diagram PNG or DOT files from existing Threat model HCL files
(as specified by <files>)
-outdir=<directory>
Directory to output PNG files. Will create directory if it doesn't exist.
Directory to output files. Will create directory if it doesn't exist.
Either this, or -out, must be set
-out=<filename>.png
Name of output PNG file. Only the first discovered data_flow_diagram will be converted into a PNG.
-out=<filename>.<png|dot>
Name of output file. Only the first discovered data_flow_diagram will be
converted. You must set the extension to png or dot depending on the mode.
Either this, or -outdir, must be set
-dot
Outputs Graphviz DOT instead. If -out or -outdir is provided files will be
generated. If neither -out or -outdir is set, then the DOT file will be
echoed to STDOUT.
Options:
-config=<file>
Expand All @@ -46,9 +53,10 @@ Options:
func (c *DfdCommand) Run(args []string) int {

flagSet := c.GetFlagset("dfd")
flagSet.StringVar(&c.flagOutDir, "outdir", "", "Directory to output PNG files. Will create directory if it doesn't exist. Either this, or -out, must be set")
flagSet.StringVar(&c.flagOutFile, "out", "", "Name of output PNG file. Either this, or -outdir, must be set")
flagSet.StringVar(&c.flagOutDir, "outdir", "", "Directory to output files. Will create directory if it doesn't exist. Either this, or -out, must be set")
flagSet.StringVar(&c.flagOutFile, "out", "", "Name of output file. Either this, or -outdir, must be set")
flagSet.BoolVar(&c.flagOverwrite, "overwrite", false, "Overwrite existing files in the outdir. Defaults to false")
flagSet.BoolVar(&c.flagDot, "dot", false, "Whether to output raw Graphviz DOT")
flagSet.Parse(args)

if c.flagConfig != "" {
Expand All @@ -60,8 +68,8 @@ func (c *DfdCommand) Run(args []string) int {
}
}

if c.flagOutDir == "" && c.flagOutFile == "" {
fmt.Printf("You must set an -outdir or -out\n\n")
if c.flagOutDir == "" && c.flagOutFile == "" && c.flagDot == false {
fmt.Printf("You must set an -outdir or -out or -dot\n\n")
fmt.Println(c.Help())
return 1
}
Expand All @@ -72,10 +80,19 @@ func (c *DfdCommand) Run(args []string) int {
return 1
}

if c.flagOutFile != "" && filepath.Ext(c.flagOutFile) != ".png" {
fmt.Printf("-out flag must end in .png\n\n")
fmt.Println(c.Help())
return 1
if c.flagDot {
if c.flagOutFile != "" && filepath.Ext(c.flagOutFile) != ".dot" {
fmt.Printf("-out flag must end in .dot\n\n")
fmt.Println(c.Help())
return 1
}
} else {

if c.flagOutFile != "" && filepath.Ext(c.flagOutFile) != ".png" {
fmt.Printf("-out flag must end in .png\n\n")
fmt.Println(c.Help())
return 1
}
}

if len(flagSet.Args()) == 0 {
Expand Down Expand Up @@ -103,7 +120,11 @@ func (c *DfdCommand) Run(args []string) int {
for _, tm := range tmParser.GetWrapped().Threatmodels {

if tm.DataFlowDiagram != nil {
outfile := outfilePath(c.flagOutDir, tm.Name, file, ".png")
fileExt := ".png"
if c.flagDot {
fileExt = ".dot"
}
outfile := outfilePath(c.flagOutDir, tm.Name, file, fileExt)

outfiles = append(outfiles, outfile)
}
Expand Down Expand Up @@ -149,15 +170,57 @@ func (c *DfdCommand) Run(args []string) int {

for _, tm := range tmParser.GetWrapped().Threatmodels {
if tm.DataFlowDiagram != nil {
if c.flagOutFile != "" {
if c.flagDot {
dot, err := tm.GenerateDot()
if err != nil {
fmt.Printf("Error generating DOT: %s\n", c.flagOutFile)
return 1
}
if c.flagOutFile != "" {
f, err := os.Create(c.flagOutFile)
if err != nil {
fmt.Printf("Error creating file %s: %s\n", c.flagOutFile, err)
return 1
}
defer f.Close()

_, err = f.WriteString(dot)
if err != nil {
fmt.Printf("Error writing DOT file to %s: %s\n", c.flagOutFile, err)
return 1
}

fmt.Printf("Successfully created '%s'\n", c.flagOutFile)
return 0

} else if c.flagOutDir != "" {
f, err := os.Create(outfilePath(c.flagOutDir, tm.Name, file, ".dot"))
if err != nil {
fmt.Printf("Error creating file %s: %s\n", outfilePath(c.flagOutDir, tm.Name, file, ".dot"), err)
return 1
}
defer f.Close()

_, err = f.WriteString(dot)
if err != nil {
fmt.Printf("Error writing DOT file to %s: %s\n", outfilePath(c.flagOutDir, tm.Name, file, ".dot"), err)
return 1
}

fmt.Printf("Successfully created '%s'\n", outfilePath(c.flagOutDir, tm.Name, file, ".dot"))
} else {
fmt.Printf("%s\n", dot)
break
}
} else if c.flagOutFile != "" {
err = tm.GenerateDfdPng(c.flagOutFile)
if err != nil {
fmt.Printf("Error generating DFD: %s\n", err)
return 1
}

fmt.Printf("Successfully created '%s'\n", c.flagOutFile)
break
return 0
} else {
err = tm.GenerateDfdPng(outfilePath(c.flagOutDir, tm.Name, file, ".png"))
if err != nil {
Expand All @@ -177,5 +240,5 @@ func (c *DfdCommand) Run(args []string) int {
}

func (c *DfdCommand) Synopsis() string {
return "Generate Data Flow Diagram PNG files from existing HCL threatmodel file(s)"
return "Generate Data Flow Diagram PNG or DOT files from existing HCL threatmodel file(s)"
}
6 changes: 3 additions & 3 deletions go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,17 +17,15 @@ require (
github.com/hashicorp/errwrap v1.1.0 // indirect
github.com/hashicorp/go-multierror v1.1.1
github.com/hashicorp/hcl/v2 v2.10.1
github.com/hashicorp/terraform-json v0.13.0 // indirect
github.com/hashicorp/terraform-json v0.13.0
github.com/imdario/mergo v0.3.12 // indirect
github.com/kami-zh/go-capturer v0.0.0-20171211120116-e492ea43421d
github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51
github.com/marqeta/go-dfd v0.1.0
github.com/mattn/go-colorable v0.1.11 // indirect
github.com/mattn/go-isatty v0.0.14 // indirect
github.com/mgutz/ansi v0.0.0-20200706080929-d51e80ef957d // indirect
github.com/microcosm-cc/bluemonday v1.0.16 // indirect
github.com/mitchellh/cli v1.1.2
github.com/mitchellh/copystructure v1.2.0 // indirect
github.com/mitchellh/go-wordwrap v1.0.1 // indirect
github.com/muesli/reflow v0.3.0 // indirect
github.com/muesli/termenv v0.9.0 // indirect
Expand All @@ -43,3 +41,5 @@ require (
golang.org/x/text v0.3.7 // indirect
gonum.org/v1/gonum v0.9.3
)

replace github.com/marqeta/go-dfd => github.com/xntrik/go-dfd v0.1.1
Loading

0 comments on commit d9f6eb2

Please sign in to comment.