refactor mikrotik formatter (#98)

* refactor mikrotik formatter * fix golangci-lint errors * embed script template as string - not via a virtual fs * refactor scriptString handling * enhance: Dont alloc a new slice to compute ipv6 just pass a template func, dont alloc a strings buffer to remove redundant new lines just alter template to not have them, pass response writer directly to the parsed template instead of allocating strings * fix: fix newlines on template * fix: remove unused struct * fix: write to temporary buffer of bytes so no partial script is written --------- Co-authored-by: Laurence Jones <[email protected]>
crowdsecurity · Jun 26, 2024 · e53c345 · e53c345
1 parent dc326b7
commit e53c345
Show file tree

Hide file tree

Showing 3 changed files with 115 additions and 34 deletions.
diff --git a/pkg/formatters/formatters.go b/pkg/formatters/formatters.go
@@ -7,12 +7,13 @@ import (
 
 	"github.com/crowdsecurity/crowdsec/pkg/models"
 
+	"github.com/crowdsecurity/cs-blocklist-mirror/pkg/formatters/mikrotik"
 	"github.com/crowdsecurity/cs-blocklist-mirror/pkg/registry"
 )
 
 var ByName = map[string]func(w http.ResponseWriter, r *http.Request){
 	"plain_text": PlainText,
-	"mikrotik":   MikroTik,
+	"mikrotik":   mikrotik.Format,
 	"f5":         F5,
 }
 
@@ -23,39 +24,6 @@ func PlainText(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-func MikroTik(w http.ResponseWriter, r *http.Request) {
-	decisions := r.Context().Value(registry.GlobalDecisionRegistry.Key).([]*models.Decision)
-
-	listName := r.URL.Query().Get("listname")
-	if listName == "" {
-		listName = "CrowdSec"
-	}
-
-	if !r.URL.Query().Has("ipv6only") {
-		fmt.Fprintf(w, "/ip firewall address-list remove [find list=%s]\n", listName)
-	}
-
-	if !r.URL.Query().Has("ipv4only") {
-		fmt.Fprintf(w, "/ipv6 firewall address-list remove [find list=%s]\n", listName)
-	}
-
-	for _, decision := range decisions {
-		var ipType = "/ip"
-		if strings.Contains(*decision.Value, ":") {
-			ipType = "/ipv6"
-		}
-
-		fmt.Fprintf(w,
-			"%s firewall address-list add list=%s address=%s comment=\"%s for %s\"\n",
-			ipType,
-			listName,
-			*decision.Value,
-			*decision.Scenario,
-			*decision.Duration,
-		)
-	}
-}
-
 func F5(w http.ResponseWriter, r *http.Request) {
 	decisions := r.Context().Value(registry.GlobalDecisionRegistry.Key).([]*models.Decision)
 	for _, decision := range decisions {

diff --git a/pkg/formatters/mikrotik/mikrotik.go b/pkg/formatters/mikrotik/mikrotik.go
@@ -0,0 +1,67 @@
+package mikrotik
+
+import (
+	"bytes"
+	_ "embed"
+	"net/http"
+	"strings"
+	"text/template"
+
+	"github.com/crowdsecurity/crowdsec/pkg/models"
+	"github.com/crowdsecurity/cs-blocklist-mirror/pkg/registry"
+)
+
+type CustomMikrotikData struct {
+	ListName               string
+	Decisions              []*models.Decision
+	NameOfMikrotikFunction string
+	IPv6Only               bool
+	IPv4Only               bool
+}
+
+//go:embed mikrotik.tmpl
+var MikrotikScriptTemplate string
+
+func Format(w http.ResponseWriter, r *http.Request) {
+
+	// Extract decisions from the context
+	decisions := r.Context().Value(registry.GlobalDecisionRegistry.Key).([]*models.Decision)
+
+	// Get query parameters
+	query := r.URL.Query()
+
+	// check if ipv6only or ipv4only is set
+	ipv6only := query.Has("ipv6only")
+	ipv4only := query.Has("ipv4only")
+
+	listName := query.Get("listname")
+	if listName == "" {
+		listName = "CrowdSec"
+	}
+
+	data := CustomMikrotikData{
+		ListName:               listName,
+		Decisions:              decisions,
+		NameOfMikrotikFunction: "CrowdSecBlockIP",
+		IPv6Only:               ipv6only,
+		IPv4Only:               ipv4only,
+	}
+
+	// Parse the template
+	parsedTemplate, err := template.New("script").Funcs(template.FuncMap{
+		"contains": strings.Contains,
+	}).Parse(MikrotikScriptTemplate)
+	if err != nil {
+		http.Error(w, "Error parsing template: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	var buf = new(bytes.Buffer)
+	// Execute the template
+	err = parsedTemplate.Execute(buf, data)
+	if err != nil {
+		http.Error(w, "Error executing template "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+	w.Write(buf.Bytes())
+}
diff --git a/pkg/formatters/mikrotik/mikrotik.tmpl b/pkg/formatters/mikrotik/mikrotik.tmpl
@@ -0,0 +1,46 @@
+{{if not $.IPv6Only -}}
+:global {{$.NameOfMikrotikFunction}} do={
+  :local list "{{$.ListName}}"
+  :local address $1
+  :local comment $2
+  :local timeout $3
+  onerror e in={ 
+    /ip firewall address-list add list=$list address=$address comment=$comment timeout="$timeout"
+  } do={
+    /ip firewall address-list remove [ find list=$list address="$address" ]
+    /ip firewall address-list add list=$list address=$address comment=$comment timeout="$timeout"
+  }
+}
+{{- if not $.IPv4Only}}
+{{end}}{{end}}
+{{- if not $.IPv4Only -}}
+:global {{$.NameOfMikrotikFunction}}v6 do={
+  :local list "{{$.ListName}}"
+  :local address $1
+  :local comment $2
+  :local timeout $3
+  onerror e in={ 
+    /ipv6 firewall address-list add list=$list address=$address comment=$comment timeout="$timeout"
+  } do={
+    /ipv6 firewall address-list remove [ find list=$list address="$address" ]
+    /ipv6 firewall address-list add list=$list address=$address comment=$comment timeout="$timeout"
+  }
+}
+{{- end -}}
+
+{{- range .Decisions}}
+{{ $ipv6Check := contains .Value ":"}}
+{{- if not $ipv6Check -}}
+${{$.NameOfMikrotikFunction}} {{.Value}} "{{.Scenario}}" {{.Duration}}
+{{- else -}}
+${{$.NameOfMikrotikFunction}}v6 {{.Value}} "{{.Scenario}}" {{.Duration}}
+{{- end }}
+{{- end }}
+
+{{- if not $.IPv6Only }}
+:set {{$.NameOfMikrotikFunction}}
+{{- end}}
+{{- if not $.IPv4Only }}
+:set {{$.NameOfMikrotikFunction}}v6
+{{- end}}
+