From 58867031ea88c8ac9f292ec8ace2678b0ac6a7a2 Mon Sep 17 00:00:00 2001
From: Alejandro Mery <amery@jpi.io>
Date: Sat, 26 Aug 2023 17:05:11 +0000
Subject: [PATCH] zones: rearrange code around WireguardKeysWriter

Signed-off-by: Alejandro Mery <amery@jpi.io>
---
 pkg/zones/machine_rings.go | 33 -----------------
 pkg/zones/wireguard.go     | 76 ++++++++++++++++++++++++++++++++++++++
 pkg/zones/zone_rings.go    | 28 --------------
 3 files changed, 76 insertions(+), 61 deletions(-)
 create mode 100644 pkg/zones/wireguard.go

diff --git a/pkg/zones/machine_rings.go b/pkg/zones/machine_rings.go
index bd67099..62c00f9 100644
--- a/pkg/zones/machine_rings.go
+++ b/pkg/zones/machine_rings.go
@@ -3,7 +3,6 @@ package zones
 import (
 	"bytes"
 	"fmt"
-	"io/fs"
 	"os"
 
 	"darvaza.org/core"
@@ -73,38 +72,6 @@ func (m *Machine) tryReadWireguardKeys(ring int) error {
 	}
 }
 
-// WriteWireguardKeys writes the wgN.key/wgN.pub files
-func (m *Machine) WriteWireguardKeys(ring int) error {
-	var err error
-	var key, pub string
-	var ri *RingInfo
-
-	ri, _ = m.getRingInfo(ring)
-	if ri != nil {
-		key = ri.Keys.PrivateKey.String()
-		pub = ri.Keys.PublicKey.String()
-	}
-
-	switch {
-	case key == "":
-		return fs.ErrNotExist
-	case pub == "":
-		pub = ri.Keys.PrivateKey.Public().String()
-	}
-
-	err = m.WriteStringFile(key+"\n", "wg%v.key", ring)
-	if err != nil {
-		return err
-	}
-
-	err = m.WriteStringFile(pub+"\n", "wg%v.pub", ring)
-	if err != nil {
-		return err
-	}
-
-	return nil
-}
-
 // RemoveWireguardKeys deletes wgN.key and wgN.pub from
 // the machine's config directory
 func (m *Machine) RemoveWireguardKeys(ring int) error {
diff --git a/pkg/zones/wireguard.go b/pkg/zones/wireguard.go
new file mode 100644
index 0000000..7fd622a
--- /dev/null
+++ b/pkg/zones/wireguard.go
@@ -0,0 +1,76 @@
+package zones
+
+import (
+	"io/fs"
+	"os"
+)
+
+var (
+	_ WireguardKeysWriter = (*Zones)(nil)
+	_ WireguardKeysWriter = (*Zone)(nil)
+	_ WireguardKeysWriter = (*Machine)(nil)
+)
+
+// A WireguardKeysWriter writes the Wireguard Keys for all machines
+// under its scope for the specified ring
+type WireguardKeysWriter interface {
+	WriteWireguardKeys(ring int) error
+}
+
+// WriteWireguardKeys rewrites all wgN.{key,pub} files
+func (m *Zones) WriteWireguardKeys(ring int) error {
+	return writeWireguardKeys(m, ring)
+}
+
+// WriteWireguardKeys rewrites all wgN.{key,pub} files on this zone
+func (z *Zone) WriteWireguardKeys(ring int) error {
+	return writeWireguardKeys(z, ring)
+}
+
+func writeWireguardKeys(m MachineIterator, ring int) error {
+	var err error
+
+	m.ForEachMachine(func(p *Machine) bool {
+		err = p.WriteWireguardKeys(ring)
+		if os.IsNotExist(err) {
+			// ignore
+			err = nil
+		}
+
+		return err != nil
+	})
+
+	return err
+}
+
+// WriteWireguardKeys writes the wgN.key/wgN.pub files
+func (m *Machine) WriteWireguardKeys(ring int) error {
+	var err error
+	var key, pub string
+	var ri *RingInfo
+
+	ri, _ = m.getRingInfo(ring)
+	if ri != nil {
+		key = ri.Keys.PrivateKey.String()
+		pub = ri.Keys.PublicKey.String()
+	}
+
+	switch {
+	case key == "":
+		return fs.ErrNotExist
+	case pub == "":
+		pub = ri.Keys.PrivateKey.Public().String()
+	}
+
+	err = m.WriteStringFile(key+"\n", "wg%v.key", ring)
+	if err != nil {
+		return err
+	}
+
+	err = m.WriteStringFile(pub+"\n", "wg%v.pub", ring)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
diff --git a/pkg/zones/zone_rings.go b/pkg/zones/zone_rings.go
index 25db1bc..bd3b5c9 100644
--- a/pkg/zones/zone_rings.go
+++ b/pkg/zones/zone_rings.go
@@ -1,29 +1,17 @@
 package zones
 
-import "os"
-
 // PruneWireguardConfig removes wgN.conf files of machines with
 // the corresponding ring disabled.
 func (z *Zone) PruneWireguardConfig(ring int) error {
 	return pruneWireguardConfig(z, ring)
 }
 
-// WriteWireguardKeys rewrites all wgN.{key,pub} files on this zone
-func (z *Zone) WriteWireguardKeys(ring int) error {
-	return writeWireguardKeys(z, ring)
-}
-
 // PruneWireguardConfig removes wgN.conf files of machines with
 // the corresponding ring disabled on all zones
 func (m *Zones) PruneWireguardConfig(ring int) error {
 	return pruneWireguardConfig(m, ring)
 }
 
-// WriteWireguardKeys rewrites all wgN.{key,pub} files
-func (m *Zones) WriteWireguardKeys(ring int) error {
-	return writeWireguardKeys(m, ring)
-}
-
 func pruneWireguardConfig(m MachineIterator, ring int) error {
 	var err error
 
@@ -37,19 +25,3 @@ func pruneWireguardConfig(m MachineIterator, ring int) error {
 
 	return err
 }
-
-func writeWireguardKeys(m MachineIterator, ring int) error {
-	var err error
-
-	m.ForEachMachine(func(p *Machine) bool {
-		err = p.WriteWireguardKeys(ring)
-		if os.IsNotExist(err) {
-			// ignore
-			err = nil
-		}
-
-		return err != nil
-	})
-
-	return err
-}