#109: Graphviz: initial code

formatter/format.go

@@ -12,13 +12,15 @@ const (
 	MarkdownOutput OutputFormat = "md"
 	// JSONOutput constant defines OutputFormat for JavaScript Object Notation, which is more useful for machine-related operations (parsing)
 	JSONOutput OutputFormat = "json"
+	// DotOutput constant defined OutputFormat for Dot (Graphviz), which can be used to generate various graphs
+	DotOutput OutputFormat = "dot"
 )
 
 // IsValid checks whether requested output format is valid
 func (of OutputFormat) IsValid() bool {
 	// markdown & md is essentially the same thing
 	switch of {
-	case "markdown", "md", "html", "csv", "json":
+	case "markdown", "md", "html", "csv", "json", "dot":
 		return true
 	}
 	return false
@@ -35,6 +37,8 @@ func (of OutputFormat) FileOutputFormat() OutputFormat {
 		return CSVOutput
 	case "json":
 		return JSONOutput
+	case "dot":
+		return DotOutput
 	}
 	return HTMLOutput
 }

formatter/formatter.go

@@ -25,6 +25,10 @@ func New(config *Config) Formatter {
 		return &CSVFormatter{
 			config,
 		}
+	case DotOutput:
+		return &DotFormatter{
+			config,
+		}
 	}
 	return nil
 }

formatter/formatter_dot.go

@@ -0,0 +1,83 @@
+package formatter
+
+import (
+	_ "embed"
+	"fmt"
+	"strings"
+	"text/template"
+)
+
+type DotFormatter struct {
+	config *Config
+}
+
+//go:embed resources/templates/graphviz.tmpl
+// DotTemplate variable is used to store contents of graphviz template
+var DotTemplate string
+
+const (
+	DotOpenPortColor     = "#228B22"
+	DotFilteredPortColor = "#FFAE00"
+	DotClosedPortColor   = "#DC143C"
+	DotDefaultColor      = "gray"
+)
+
+// Format the data and output it to appropriate io.Writer
+func (f *DotFormatter) Format(td *TemplateData, templateContent string) (err error) {
+	tmpl := template.New("dot")
+	f.defineTemplateFunctions(tmpl)
+	tmpl, err = tmpl.Parse(templateContent)
+	if err != nil {
+		return
+	}
+	return tmpl.Execute(f.config.Writer, td)
+}
+
+// defaultTemplateContent returns default template content for any typical chosen formatter (HTML or Markdown)
+func (f *DotFormatter) defaultTemplateContent() string {
+	return DotTemplate
+}
+
+// defineTemplateFunctions defines all template functions that are used in dot templates
+func (f *DotFormatter) defineTemplateFunctions(tmpl *template.Template) {
+	tmpl.Funcs(
+		template.FuncMap{
+			"clean_ip":         cleanIP,
+			"port_state_color": portStateColor,
+			"hop_list":         hopList,
+		},
+	)
+}
+
+// cleanIP removes dots from IP address to make it possible to use in graphviz as an ID
+func cleanIP(ip string) string {
+	return strings.ReplaceAll(ip, ".", "")
+}
+
+// portStateColor returns hexademical color value for state port
+func portStateColor(port *Port) string {
+	switch port.State.State {
+	case "open":
+		return DotOpenPortColor
+	case "filtered":
+		return DotFilteredPortColor
+	case "closed":
+		return DotClosedPortColor
+	}
+	return DotDefaultColor
+}
+
+func hopList(hops []Hop, startHop string, endHopName string, endHopKey int) map[string]string {
+	var hopList map[string]string = map[string]string{}
+	var previous *Hop = nil
+	for i := range hops {
+		if i == 0 {
+			hopList[startHop] = fmt.Sprintf("hop%s", hops[i].IPAddr)
+		} else {
+			hopList[fmt.Sprintf("hop%s", previous.IPAddr)] = fmt.Sprintf("hop%s", hops[i].IPAddr)
+		}
+		previous = &hops[i]
+	}
+	hopList[fmt.Sprintf("hop%s", previous.IPAddr)] = fmt.Sprintf("%s%d", endHopName, endHopKey)
+	return hopList
+}

formatter/formatter_test.go

@@ -61,6 +61,15 @@ func TestNew(t *testing.T) {
 			},
 			want: &CSVFormatter{config: &Config{OutputFormat: CSVOutput}},
 		},
+		{
+			name: "DOT output",
+			args: args{
+				config: &Config{
+					OutputFormat: DotOutput,
+				},
+			},
+			want: &DotFormatter{config: &Config{OutputFormat: DotOutput}},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

formatter/nmap_host.go

@@ -1,6 +1,10 @@
 package formatter
 
-import "fmt"
+import (
+	"encoding/xml"
+	"fmt"
+	"strconv"
+)
 
 // Host describes host related entry (`host` node)
 type Host struct {
@@ -126,8 +130,20 @@ type Trace struct {
 
 // Hop struct contains information about HOP record with time to live, host name, IP
 type Hop struct {
-	TTL    int     `xml:"ttl,attr"`
-	IPAddr string  `xml:"ipaddr,attr"`
-	RTT    float64 `xml:"rtt,attr"`
-	Host   string  `xml:"host,attr"`
+	TTL    int    `xml:"ttl,attr"`
+	IPAddr string `xml:"ipaddr,attr"`
+	RTT    RTT    `xml:"rtt,attr"`
+	Host   string `xml:"host,attr"`
+}
+
+type RTT float64
+
+// UnmarshalXMLAttr
+func (r *RTT) UnmarshalXMLAttr(attr xml.Attr) error {
+	value, err := strconv.ParseFloat(attr.Value, 64)
+	if err != nil {
+		value = 0.0
+	}
+	*(*float64)(r) = value
+	return nil
 }

formatter/nmap_host_test.go

@@ -1,6 +1,8 @@
 package formatter
 
-import "testing"
+import (
+	"testing"
+)
 
 func TestHost_JoinedAddresses(t *testing.T) {
 	type fields struct {

formatter/nmap_run.go

@@ -55,3 +55,16 @@ type StatHosts struct {
 	Down  int `xml:"down,attr"`
 	Total int `xml:"total,attr"`
 }
+
+// AllHops is getting all possible hops that occurred during the scan and
+// merges them uniquely into one map
+func (n *NMAPRun) AllHops() map[string]Hop {
+	var hops map[string]Hop = map[string]Hop{}
+	for i := range n.Host {
+		for j := range n.Host[i].Trace.Hops {
+			hop := n.Host[i].Trace.Hops[j]
+			hops[hop.IPAddr] = hop
+		}
+	}
+	return hops
+}

formatter/resources/templates/graphviz.tmpl

@@ -0,0 +1,29 @@
+strict digraph "Scan" {
+    fontname="monospace"
+    node [fontname="monospace", width=.25, height=.375, fontsize=9]
+    edge [fontname="monospace"]
+    bgcolor="#111"
+    layout=dot
+
+    scanner [label="{{ .NMAPRun.Scanner }}", shape=hexagon, style=filled];
+
+    {{- $hops := .NMAPRun.AllHops }}
+    {{ range $key, $value := $hops -}}
+    hop{{ clean_ip $key }} [label="", tooltip="{{ $key }}", shape=circle, height=.12, width=.12, style=filled]
+    {{ end }}
+
+    {{ range $key, $value := .NMAPRun.Host -}}
+    {{ if eq $value.Status.State "up" -}}
+    srv{{ $key }} [label="{{ $value.JoinedAddresses "/" }}" tooltip="{{ $value.JoinedAddresses "/" }}", shape=hexagon, style=filled];
+    {{ range $portKey, $portValue := $value.Port }}
+    srv{{ $key }}_port_{{ $portValue.Protocol }}_{{ $portValue.PortID }} [label="{{ $portValue.Protocol }}/{{ $portValue.PortID }} ({{ $portValue.State.State }})", tooltip="{{ $portValue.Protocol }}/{{ $portValue.PortID }} ({{ $portValue.State.State }})", shape=underline, width=.12, height=.12, fontsize=8, color="{{ port_state_color $portValue }}", fontcolor=gray];
+    srv{{ $key }} -> srv{{ $key }}_port_{{ $portValue.Protocol }}_{{ $portValue.PortID }} [arrowhead=none, color=gray];
+    {{ end -}}
+
+    {{ range $hopSource, $hopTarget := hop_list $value.Trace.Hops "scanner" "srv" $key }}
+    {{ clean_ip $hopSource }} -> {{ clean_ip $hopTarget }} [arrowhead=none, color=gray]
+    {{- end }}
+
+    {{- end }}
+    {{ end }}
+}