netbox: add GetDnsServersForSiteclients/netbox/v4.1.0

3 files changed, 91 insertions, 0 deletions

diff --git a/clients/netbox/client.go b/clients/netbox/client.go
index 3b7fbaf..aceebed 100644
--- a/clients/netbox/client.go
+++ b/clients/netbox/client.go
@@ -11,6 +11,7 @@ type NetboxClient interface {
         GetSitePrefixesWithTag(ctx context.Context, site string, tag string) ([]*net.IPNet, error)
         GetPrefixesWithTag(ctx context.Context, tag string) ([]*net.IPNet, error)
         GetServicesForVm(ctx context.Context, vmName string) ([]*Service, error)
+        GetDnsServersForSite(ctx context.Context, site string) ([]net.IP, error)
 }
 
 type BasicNetboxClient struct {
@@ -145,3 +146,25 @@ func (c *BasicNetboxClient) GetServicesForVm(ctx context.Context, vmName string)
 
         return out, nil
 }
+
+func (c *BasicNetboxClient) GetDnsServersForSite(ctx context.Context, site string) ([]net.IP, error) {
+        var m dnsServerGQLResponse
+        if err := c.Do(ctx, &NetboxGraphQLRequest{dnsServerGQLQuery}, &m); err != nil {
+                return nil, err
+        }
+
+        out := []net.IP{}
+        for _, a := range m.Data.AddressList {
+                if a.AssignedObject.VirtualMachine.Site.Name == site ||
+                        a.AssignedObject.Device.Site.Name == site {
+
+                        ip, _, err := net.ParseCIDR(a.Address)
+                        if err != nil {
+                                return nil, err
+                        }
+
+                        out = append(out, ip)
+                }
+        }
+        return out, nil
+}
diff --git a/clients/netbox/config_file_client.go b/clients/netbox/config_file_client.go
index d84dd80..c692b9f 100644
--- a/clients/netbox/config_file_client.go
+++ b/clients/netbox/config_file_client.go
@@ -150,6 +150,18 @@ func parseValues(v []string) ([]*net.IPNet, error) {
         return out, nil
 }
 
+func parseIps(v []string) ([]net.IP, error) {
+        out := make([]net.IP, len(v))
+        for i, n := range v {
+                ip := net.ParseIP(n)
+                if ip == nil {
+                        return nil, fmt.Errorf("Invalid IP '%s'", n)
+                }
+                out[i] = ip
+        }
+        return out, nil
+}
+
 // GetServicesForVm returns a list of services for a named VM
 //
 // The data format of the config file should be, under the key "services":
@@ -183,3 +195,47 @@ func (c *ConfigFileNetboxClient) GetServicesForVm(ctx context.Context, vmName st
 
         return vm, nil
 }
+
+// GetDnsServersForSite gets a list of IP address for a site that are
+// assigned to DNS servers.
+//
+// The data format of the config file should be, under the key "sites":
+//
+//      map[string]map[string][]string{
+//              "site_name": map[string][]string{
+//                      "dns-servers": []string{
+//                              "127.0.0.1",
+//                      },
+//              },
+//      }
+//
+// In JSON format:
+//
+//      { "sites": {
+//          "site": {
+//            "dns-servers": [ "127.0.0.1" ]
+//          }
+//      }
+func (c *ConfigFileNetboxClient) GetDnsServersForSite(ctx context.Context, site string) ([]net.IP, error) {
+        s, ok := c.cfg["sites"]
+        if !ok {
+                return nil, fmt.Errorf("No key 'sites' in config file")
+        }
+
+        sites := map[string]map[string][]string{}
+        if err := mapstructure.Decode(s, &sites); err != nil {
+                return nil, fmt.Errorf("Key 'sites' not in correct format: %w", err)
+        }
+
+        tags, ok := sites[site]
+        if !ok {
+                return nil, fmt.Errorf("Site %s not in sites", site)
+        }
+
+        values, ok := tags["dns-servers"]
+        if !ok {
+                return nil, fmt.Errorf("Site %s not in sites", site)
+        }
+
+        return parseIps(values)
+}
diff --git a/clients/netbox/config_file_client_test.go b/clients/netbox/config_file_client_test.go
index 117db00..7af4647 100644
--- a/clients/netbox/config_file_client_test.go
+++ b/clients/netbox/config_file_client_test.go
@@ -13,11 +13,15 @@ import (
 var (
         ipNet1 *net.IPNet
         ipNet2 *net.IPNet
+        ip1    net.IP
+        ip2    net.IP
 )
 
 func init() {
         _, ipNet1, _ = net.ParseCIDR("127.0.0.1/8")
         _, ipNet2, _ = net.ParseCIDR("10.0.10.0/24")
+        ip1 = net.ParseIP("127.0.0.1")
+        ip2 = net.ParseIP("127.0.0.2")
 }
 
 var testFs fstest.MapFS
@@ -57,6 +61,7 @@ netbox:
   sites:
     site1:
       tag1: [ "127.0.0.1/8", "10.0.10.0/24" ]
+      dns-servers: [ "127.0.0.1", "127.0.0.2" ]
   tags:
     tag1: [ "127.0.0.1/8", "10.0.10.0/24" ]
   services:
@@ -186,6 +191,13 @@ func (s *ConfigFileNetboxClientSuite) TestGetServicesForVMMissingVM() {
         assert.ErrorContains(s.T(), err, "VM named foo not in config")
 }
 
+func (s *ConfigFileNetboxClientSuite) TestGetDnsServersForSite() {
+        ips, err := s.c.GetDnsServersForSite(context.TODO(), "site1")
+        assert.NoError(s.T(), err)
+        assert.True(s.T(), ip1.Equal(ips[0]))
+        assert.True(s.T(), ip2.Equal(ips[1]))
+}
+
 func TestConfigFileNetboxClientSuite(t *testing.T) {
         suite.Run(t, &ConfigFileNetboxClientSuite{})
 }