feat: d2-lang support

Author: vdjagilev

cmd/root.go

@@ -52,9 +52,9 @@ var workflow formatter.Workflow
 
 // rootCmd represents the base command when called without any subcommands
 var rootCmd = &cobra.Command{
-	Use:   "nmap-formatter [html|csv|md|json|dot|sqlite|excel] [path-to-nmap.xml]",
+	Use:   "nmap-formatter [html|csv|md|json|dot|sqlite|excel|d2] [path-to-nmap.xml]",
 	Short: "Utility that can help you to convert NMAP XML application output to various other formats",
-	Long:  `This utility allows you to convert NMAP XML output to various other formats like (html, csv, markdown (md), json, dot, excel, sqlite)`,
+	Long:  `This utility allows you to convert NMAP XML output to various other formats like (html, csv, markdown (md), json, dot, excel, sqlite, d2)`,
 	Args:  arguments,
 	RunE:  run,
 }
@@ -114,6 +114,9 @@ func init() {
 	rootCmd.Flags().StringVar(&config.OutputOptions.SqliteOutputOptions.DSN, "sqlite-dsn", "nmap.sqlite", "--sqlite-dsn nmap.sqlite")
 	rootCmd.Flags().StringVar(&config.OutputOptions.SqliteOutputOptions.ScanIdentifier, "scan-id", "", "--scan-id abc123")
 
+	// Configs related to D2 language
+	rootCmd.Flags().BoolVar(&config.OutputOptions.D2LangOptions.SkipDownHosts, "d2-skip-down-hosts", true, "--d2-skip-down-hosts=false, would print all hosts that are offline in D2 language output")
+
 	workflow = &formatter.MainWorkflow{}
 }
 

formatter/format.go

@@ -18,13 +18,15 @@ const (
 	SqliteOutput OutputFormat = "sqlite"
 	// ExcelOutput constant defines OutputFormat for Excel file, which can be used to generate Excel files
 	ExcelOutput OutputFormat = "excel"
+	// D2LangOutput constant defines OutputFormat for D2 language, which can be used to generate D2 language files
+	D2LangOutput OutputFormat = "d2"
 )
 
 // 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", "dot", "sqlite", "excel":
+	case "markdown", "md", "html", "csv", "json", "dot", "sqlite", "excel", "d2":
 		return true
 	}
 	return false

formatter/formatter.go

@@ -37,6 +37,10 @@ func New(config *Config) Formatter {
 		return &SqliteFormatter{
 			config,
 		}
+	case D2LangOutput:
+		return &D2LangFormatter{
+			config,
+		}
 	}
 	return nil
 }

formatter/formatter_d2.go

@@ -0,0 +1,78 @@
+package formatter
+
+import (
+	"context"
+	"encoding/hex"
+	"fmt"
+	"hash/fnv"
+
+	"oss.terrastruct.com/d2/d2format"
+	"oss.terrastruct.com/d2/d2graph"
+	"oss.terrastruct.com/d2/d2layouts/d2dagrelayout"
+	"oss.terrastruct.com/d2/d2lib"
+	"oss.terrastruct.com/d2/d2oracle"
+	"oss.terrastruct.com/d2/lib/log"
+	"oss.terrastruct.com/d2/lib/textmeasure"
+)
+
+type D2LangFormatter struct {
+	config *Config
+}
+
+func (f *D2LangFormatter) Format(td *TemplateData, templateContent string) (err error) {
+	ruler, _ := textmeasure.NewRuler()
+	layoutResolver := func(engine string) (d2graph.LayoutGraph, error) {
+		return d2dagrelayout.DefaultLayout, nil
+	}
+	compileOpts := &d2lib.CompileOptions{
+		LayoutResolver: layoutResolver,
+		Ruler:          ruler,
+	}
+
+	_, graph, _ := d2lib.Compile(log.Stderr(context.Background()), "nmap", compileOpts, nil)
+
+	for i := range td.NMAPRun.Host {
+		host := &td.NMAPRun.Host[i]
+		fnv := fnv.New128()
+
+		if host.ShouldSkipHost(td.OutputOptions.D2LangOptions.SkipDownHosts) {
+			continue
+		}
+
+		address := host.JoinedAddresses("/")
+		hostnames := host.JoinedHostNames("/")
+		hostLabel := address
+		if hostnames != "" {
+			hostLabel = fmt.Sprintf("%s\n(%s)", address, hostnames)
+		}
+		_, err := fnv.Write([]byte(address))
+		if err != nil {
+			return err
+		}
+
+		hostId := hex.EncodeToString(fnv.Sum(nil))
+		graph, _, _ = d2oracle.Create(graph, nil, hostId)
+		graph, _ = d2oracle.Set(graph, nil, hostId+".label", nil, &hostLabel)
+		graph, _ = d2oracle.Set(graph, nil, "nmap -> "+hostId, nil, nil)
+
+		for j := range host.Port {
+			port := &host.Port[j]
+			portId := fmt.Sprintf("%s-port%d", hostId, port.PortID)
+			graph, _, _ = d2oracle.Create(graph, nil, portId)
+			portLabel := fmt.Sprintf("%d/%s\n%s\n%s", port.PortID, port.Protocol, port.State.State, port.Service.Name)
+			graph, _ = d2oracle.Set(graph, nil, portId+".label", nil, &portLabel)
+			shape := "circle"
+			graph, _ = d2oracle.Set(graph, nil, portId+".shape", nil, &shape)
+			width := "25"
+			graph, _ = d2oracle.Set(graph, nil, portId+".width", nil, &width)
+			graph, _ = d2oracle.Move(graph, nil, portId, hostId+"."+portId, true)
+		}
+	}
+	_, err = f.config.Writer.Write([]byte(d2format.Format(graph.AST)))
+	return
+}
+
+// defaultTemplateContent does not return anything in this case
+func (f *D2LangFormatter) defaultTemplateContent() string {
+	return ""
+}

formatter/output_options.go

@@ -8,6 +8,7 @@ type OutputOptions struct {
 	CSVOptions          CSVOutputOptions
 	SqliteOutputOptions SqliteOutputOptions
 	ExcelOptions        ExcelOutputOptions
+	D2LangOptions       D2LangOutputOptions
 }
 
 // HTMLOutputOptions stores options related only to HTML conversion/formatting
@@ -68,3 +69,9 @@ type ExcelOutputOptions struct {
 	// The hosts that are down won't be displayed
 	SkipDownHosts bool
 }
+
+// D2LangOutputOptions store options related to D2 language file formatting
+type D2LangOutputOptions struct {
+	// The hosts that are down won't be displayed
+	SkipDownHosts bool
+}

go.mod

@@ -7,16 +7,43 @@ require (
 	github.com/mattn/go-sqlite3 v1.14.22
 	github.com/spf13/cobra v1.8.0
 	golang.org/x/net v0.24.0
+	oss.terrastruct.com/d2 v0.6.3
 )
 
 require (
+	cdr.dev/slog v1.4.2-0.20221206192828-e4803b10ae17 // indirect
+	github.com/PuerkitoBio/goquery v1.8.1 // indirect
+	github.com/alecthomas/chroma v0.10.0 // indirect
+	github.com/alecthomas/chroma/v2 v2.5.0 // indirect
+	github.com/andybalholm/cascadia v1.3.2 // indirect
+	github.com/dlclark/regexp2 v1.10.0 // indirect
+	github.com/dop251/goja v0.0.0-20231027120936-b396bb4c349d // indirect
+	github.com/fatih/color v1.13.0 // indirect
+	github.com/go-sourcemap/sourcemap v2.1.3+incompatible // indirect
+	github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0 // indirect
+	github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e // indirect
+	github.com/google/pprof v0.0.0-20231205033806-a5a03c77bf08 // indirect
+	github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
+	github.com/mattn/go-colorable v0.1.9 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/mazznoer/csscolorparser v0.1.3 // indirect
 	github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 // indirect
 	github.com/richardlehane/mscfb v1.0.4 // indirect
 	github.com/richardlehane/msoleps v1.0.3 // indirect
+	github.com/rivo/uniseg v0.4.4 // indirect
 	github.com/xuri/efp v0.0.0-20231025114914-d1ff6096ae53 // indirect
 	github.com/xuri/nfp v0.0.0-20240318013403-ab9948c2c4a7 // indirect
+	github.com/yuin/goldmark v1.6.0 // indirect
+	go.opencensus.io v0.24.0 // indirect
 	golang.org/x/crypto v0.22.0 // indirect
+	golang.org/x/exp v0.0.0-20231127185646-65229373498e // indirect
+	golang.org/x/image v0.14.0 // indirect
+	golang.org/x/sys v0.19.0 // indirect
+	golang.org/x/term v0.19.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
+	golang.org/x/xerrors v0.0.0-20231012003039-104605ab7028 // indirect
+	gonum.org/v1/plot v0.14.0 // indirect
+	oss.terrastruct.com/util-go v0.0.0-20231101220827-55b3812542c2 // indirect
 )
 
 require (