From 4a4105de661955cf84b62253c4038201a7526e44 Mon Sep 17 00:00:00 2001
From: Mike Crute <mike@crute.us>
Date: Wed, 27 Sep 2023 16:56:29 -0700
Subject: netbox: add GetDnsServersForSite

---
 clients/netbox/client.go                  | 23 +++++++++++++
 clients/netbox/config_file_client.go      | 56 +++++++++++++++++++++++++++++++
 clients/netbox/config_file_client_test.go | 12 +++++++
 3 files changed, 91 insertions(+)

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{})
 }
-- 
cgit v1.2.3