From 6d05f0501f56e6c16c1aea3397aa06865a8a6b4e Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Thu, 2 Jan 2025 18:16:24 +0100
Subject: [PATCH] network policy now fully generated

---
 cmd/policygen/config.go                     | 15 ++++--
 cmd/policygen/main.go                       |  1 -
 cmd/policygen/templates.go                  |  2 +-
 cmd/policygen/templates/netpol/pod/pod.yaml | 60 ++++++++++++---------
 example/config.yaml                         | 11 ++--
 5 files changed, 55 insertions(+), 34 deletions(-)

diff --git a/cmd/policygen/config.go b/cmd/policygen/config.go
index da52a97..caa1467 100644
--- a/cmd/policygen/config.go
+++ b/cmd/policygen/config.go
@@ -39,17 +39,22 @@ func (c CIDR) MarshalYAML() ([]byte, error) {
 	return []byte(string(c)), nil
 }
 
+type Port struct {
+	Port     string `yaml:"port"`
+	Protocol string `yaml:"protocol"`
+}
+
 // Network represents each network entry in the YAML
 type Network struct {
-	Name   string   `yaml:"name"`
-	CIDR   CIDR     `yaml:"cidr"`
-	Except []CIDR   `yaml:"except,omitempty"`
-	Ports  []string `yaml:"ports,omitempty"`
+	Name   string `yaml:"name"`
+	CIDR   CIDR   `yaml:"cidr"`
+	Except []CIDR `yaml:"except,omitempty"`
+	Ports  []Port `yaml:"ports,omitempty"`
 }
 
 type Application struct {
 	Name        string            `yaml:"name"`
-	Ports       []string          `yaml:"ports,omitempty"`
+	Ports       []Port            `yaml:"ports,omitempty"`
 	MatchLabels map[string]string `yaml:"matchLabels"`
 	Namespace   string            `yaml:"-"`
 }
diff --git a/cmd/policygen/main.go b/cmd/policygen/main.go
index 8f9f7fc..dcaeb69 100644
--- a/cmd/policygen/main.go
+++ b/cmd/policygen/main.go
@@ -22,7 +22,6 @@ func execute(files []string, options *Options) error {
 		if err != nil {
 			return err
 		}
-		fmt.Printf("PARSED %+v\n", config)
 
 		policyTemplates, err := NewPolicyTemplates()
 		if err != nil {
diff --git a/cmd/policygen/templates.go b/cmd/policygen/templates.go
index 9ce5b61..80de7c6 100644
--- a/cmd/policygen/templates.go
+++ b/cmd/policygen/templates.go
@@ -55,7 +55,7 @@ func showContents(files fs.FS) {
 		panic(err)
 	}
 	for _, entry := range entries {
-		fmt.Printf("entry %s %s\n", entry.Name(), entry.Type())
+		fmt.Fprintf(os.Stderr, "entry %s %s\n", entry.Name(), entry.Type())
 		if entry.Type().IsDir() {
 			subdir, err := fs.Sub(files, entry.Name())
 			if err != nil {
diff --git a/cmd/policygen/templates/netpol/pod/pod.yaml b/cmd/policygen/templates/netpol/pod/pod.yaml
index c606036..d9ee62d 100644
--- a/cmd/policygen/templates/netpol/pod/pod.yaml
+++ b/cmd/policygen/templates/netpol/pod/pod.yaml
@@ -1,5 +1,14 @@
 ---
-{{- define "peer" }}
+{{- define "ports" }}
+          {{- range $port := . }}
+          - port: {{ $port.Port }}
+            {{- if $port.Protocol }}
+            protocol: {{ $port.Protocol }}
+            {{- end }}
+          {{- end }}
+{{- end }}
+{{- define "peers" }}
+      {{- range .Applications }}
       - podSelector:
           matchLabels: {{ .MatchLabels | toYaml | nindent 12 }}
         namespaceSelector:
@@ -7,19 +16,24 @@
             kubernetes.io/metadata.name: {{ .Namespace }}
         {{- if .Ports }}
         ports:
-          # TODO: add protocol
-          {{- range $port := .Ports }}
-          - port: {{ $port }}
-          {{- end }}
+          {{- template "ports" .Ports }}
         {{- end }}
+      {{- end }}
+{{- end }}
+{{- define "networks" }}
+      {{- range .Networks }}
+      - ipBlock:
+          cidr: {{ .CIDR}}
+          except:
+          {{- range $except := .Except }}
+          - {{ $except }}
+          {{- end }}
+        {{- if .Ports }}
+        ports:
+          {{- template "ports" .Ports }}
+        {{- end }}
+      {{- end }}
 {{- end }}
--
-  {{- define "ports" }}
-    {{- range $port := . }}
-    PORT {{ $port }}
-    {{- end }}
-  {{- end }}
-
 
 kind: NetworkPolicy
 apiVersion: networking.k8s.io/v1
@@ -40,17 +54,15 @@ spec:
   {{- if or .ingress.Applications .ingress.Networks }}
   ingress:
     from:
-      {{- range $ingress := .ingress.Applications }}
-        {{- template "peer" $ingress }}
-        {{- template "ports" $ingress.Ports }}
-      {{- end }}
-      {{- range $ingress := .ingress.Networks }}
-      - ipBlock:
-          cidr: {{ $ingress.CIDR}}
-          except:
-          {{- range $except := $ingress.Except }}
-          - {{ $except }}
-          {{- end }}
-      {{- end }}
+      {{- template "peers" .ingress }}
+      {{- template "networks" .ingress }}
   {{- end }}
+
+  {{- if or .egress.Applications .egress.Networks }}
+  egress:
+    tp:
+      {{- template "peers" .egress }}
+      {{- template "networks" .egress }}
+  {{- end }}
+
     THEEND
\ No newline at end of file
diff --git a/example/config.yaml b/example/config.yaml
index a5cc0a6..166225c 100644
--- a/example/config.yaml
+++ b/example/config.yaml
@@ -8,6 +8,8 @@ networks:
       - 10.0.0.0/8
       - 172.16.0.0/12
       - 192.168.0.0/16
+    ports:
+      - port: 2303
 
 
 namespaces:
@@ -20,8 +22,9 @@ namespaces:
         # ports when specified at the application level are used when
         # not explicitly mentioned when a link is made
         ports:
-          - 8081
-          - 8082
+          - port: 8081
+          - port: 8082
+            protocol: UDP
         matchLabels:
           app: nexus-server
 
@@ -31,7 +34,9 @@ namespaces:
         matchLabels:
           app: wamblee-org
         ports:
-         - 1000
+         - port: 1000
+         - port: 1001
+           protocol: UDP
 
 communications:
   - from: # can we support both string and list of strings?