netbox: add GetServicesForVm callclients/netbox/v3.3.0

4 files changed, 179 insertions, 18 deletions

diff --git a/clients/netbox/client.go b/clients/netbox/client.go
index 0cad733..ad2967a 100644
--- a/clients/netbox/client.go
+++ b/clients/netbox/client.go
@@ -14,6 +14,7 @@ import (
 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)
 }
 
 type BasicNetboxClient struct {
@@ -157,3 +158,32 @@ func (c *BasicNetboxClient) resolveSiteNameToId(ctx context.Context, s string) (
 
         return out.Results[0].ID, nil
 }
+
+func (c *BasicNetboxClient) GetServicesForVm(ctx context.Context, vmName string) ([]*Service, error) {
+        out := []*Service{}
+
+        q := url.Values{}
+        q.Add("virtual_machine", vmName)
+
+        page := &ServiceList{}
+        if err := c.makeRequest(ctx, http.MethodGet, "/api/ipam/services/", q, nil, page); err != nil {
+                return nil, err
+        }
+
+        for _, r := range page.Results {
+                out = append(out, r)
+        }
+
+        for page.Next != "" {
+                page = &ServiceList{}
+                if err := c.makeRequestRaw(ctx, http.MethodGet, page.Next, nil, page); err != nil {
+                        return nil, err
+                }
+
+                for _, r := range page.Results {
+                        out = append(out, r)
+                }
+        }
+
+        return out, nil
+}
diff --git a/clients/netbox/config_file_client.go b/clients/netbox/config_file_client.go
index 6023916..d84dd80 100644
--- a/clients/netbox/config_file_client.go
+++ b/clients/netbox/config_file_client.go
@@ -67,21 +67,21 @@ func NewConfigFileClient(filesystem fs.FS, name, key string) (NetboxClient, erro
 //
 // The data format of the config file should be, under the key "sites":
 //
-//  map[string]map[string][]string{
-//      "site_name": map[string][]string{
-//              "tag_name": []string{
-//                      "127.0.0.1/24",
-//              },
-//      },
-//  }
+//      map[string]map[string][]string{
+//              "site_name": map[string][]string{
+//                      "tag_name": []string{
+//                              "127.0.0.1/24",
+//                      },
+//              },
+//      }
 //
 // In JSON format:
 //
-// { "sites": {
-//     "site": {
-//       "tag_name": [ "127.0.0.1/24" ]
-//     }
-// }
+//      { "sites": {
+//          "site": {
+//            "tag_name": [ "127.0.0.1/24" ]
+//          }
+//      }
 func (c *ConfigFileNetboxClient) GetSitePrefixesWithTag(ctx context.Context, site string, tag string) ([]*net.IPNet, error) {
         s, ok := c.cfg["sites"]
         if !ok {
@@ -110,11 +110,11 @@ func (c *ConfigFileNetboxClient) GetSitePrefixesWithTag(ctx context.Context, sit
 //
 // The data format of the config file should be, under the key "tags":
 //
-//  map[string][]string{
-//      "tag_name": []string{
-//          "127.0.0.1/24",
-//      },
-//  }
+//      map[string][]string{
+//              "tag_name": []string{
+//                  "127.0.0.1/24",
+//              },
+//      }
 //
 // In JSON format:
 //
@@ -149,3 +149,37 @@ func parseValues(v []string) ([]*net.IPNet, error) {
         }
         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":
+//
+//      map[string]map[string][]string{
+//              "vm_name": []*Service{},
+//      }
+//
+// In JSON format:
+//
+//      { "services": {
+//          "myVm": [
+//            { "id": "myid }
+//          ]
+//      }
+func (c *ConfigFileNetboxClient) GetServicesForVm(ctx context.Context, vmName string) ([]*Service, error) {
+        s, ok := c.cfg["services"]
+        if !ok {
+                return nil, fmt.Errorf("No key 'services' in config file")
+        }
+
+        services := map[string][]*Service{}
+        if err := mapstructure.Decode(s, &services); err != nil {
+                return nil, fmt.Errorf("Key 'services' not in correct format: %w", err)
+        }
+
+        vm, ok := services[vmName]
+        if !ok {
+                return nil, fmt.Errorf("VM named %s not in config", vmName)
+        }
+
+        return vm, nil
+}
diff --git a/clients/netbox/config_file_client_test.go b/clients/netbox/config_file_client_test.go
index bce5165..117db00 100644
--- a/clients/netbox/config_file_client_test.go
+++ b/clients/netbox/config_file_client_test.go
@@ -59,8 +59,20 @@ netbox:
       tag1: [ "127.0.0.1/8", "10.0.10.0/24" ]
   tags:
     tag1: [ "127.0.0.1/8", "10.0.10.0/24" ]
+  services:
+    testVM:
+      - id: 1
+        name: Test Service
+        virtual_machine:
+          id: 2
+          name: Test VM
 `)
 
+var serviceTestConfig = []byte(`{
+        "netbox": {
+        }
+}`)
+
 func init() {
         testFs = fstest.MapFS{}
         testFs["netbox.foo"] = &fstest.MapFile{Data: jsonTestConfig}
@@ -160,6 +172,20 @@ func (s *ConfigFileNetboxClientSuite) TestGetPrefixesWithTagMissingInvalid() {
         assert.ErrorContains(s.T(), err, "Key 'tags' not in correct format")
 }
 
+func (s *ConfigFileNetboxClientSuite) TestGetServicesForVM() {
+        svcs, err := s.c.GetServicesForVm(context.TODO(), "testVM")
+        assert.NoError(s.T(), err)
+        assert.Equal(s.T(), svcs[0].ID, 1)
+        assert.Equal(s.T(), svcs[0].Name, "Test Service")
+        assert.Equal(s.T(), svcs[0].VirtualMachine.ID, 2)
+        assert.Equal(s.T(), svcs[0].VirtualMachine.Name, "Test VM")
+}
+
+func (s *ConfigFileNetboxClientSuite) TestGetServicesForVMMissingVM() {
+        _, err := s.c.GetServicesForVm(context.TODO(), "foo")
+        assert.ErrorContains(s.T(), err, "VM named foo not in config")
+}
+
 func TestConfigFileNetboxClientSuite(t *testing.T) {
         suite.Run(t, &ConfigFileNetboxClientSuite{})
 }
diff --git a/clients/netbox/model.go b/clients/netbox/model.go
index 40c4aa6..be4767d 100644
--- a/clients/netbox/model.go
+++ b/clients/netbox/model.go
@@ -1,6 +1,10 @@
 package netbox
 
-import "fmt"
+import (
+        "encoding/json"
+        "fmt"
+        "net"
+)
 
 type ApiError struct {
         Status int
@@ -85,3 +89,70 @@ type Prefix struct {
         Tags         []*Tag                 `json:"tags"`
         CustomFields map[string]interface{} `json:"custom_fields"`
 }
+
+type ServiceList struct {
+        Count    int        `json:"count"`
+        Next     string     `json:"next"`
+        Previous string     `json:"previous"`
+        Results  []*Service `json:"results"`
+}
+
+type ServiceProtocol struct {
+        Value string `json:"value"`
+        Label string `json:"label"`
+}
+
+type DeviceOrVM struct {
+        ID      int    `json:"id"`
+        Url     string `json:"url"`
+        Name    string `json:"name"`
+        Display string `json:"display"`
+}
+
+type ServiceIPAddress struct {
+        ID      int    `json:"id"`
+        Url     string `json:"url"`
+        Display string `json:"display"`
+        Family  int    `json:"family"`
+        Address net.IP
+        Network *net.IPNet
+}
+
+func (a *ServiceIPAddress) UnmarshalJSON(d []byte) error {
+        type Alias ServiceIPAddress
+        v := &struct {
+                *Alias
+                Address string `json:"address"`
+        }{Alias: (*Alias)(a)}
+
+        if err := json.Unmarshal(d, v); err != nil {
+                return err
+        }
+
+        i, n, err := net.ParseCIDR(v.Address)
+        if err != nil {
+                return err
+        }
+
+        a.Address = i
+        a.Network = n
+
+        return nil
+}
+
+type Service struct {
+        ID             int                    `json:"id" mapstructure:"id"`
+        Url            string                 `json:"url" mapstructure:"url"`
+        Name           string                 `json:"name" mapstructure:"name"`
+        Display        string                 `json:"display" mapstructure:"display"`
+        Device         *DeviceOrVM            `json:"device" mapstructure:"device"`
+        VirtualMachine *DeviceOrVM            `json:"virtual_machine" mapstructure:"virtual_machine"`
+        Ports          []int                  `json:"ports" mapstructure:"ports"`
+        Protocol       *ServiceProtocol       `json:"protocol" mapstructure:"protocol"`
+        Addresses      []*ServiceIPAddress    `json:"ipaddresses" mapstructure:"ipaddresses"`
+        Description    string                 `json:"description" mapstructure:"description"`
+        Created        string                 `json:"created" mapstructure:"created"`
+        LastUpdated    string                 `json:"last_updated" mapstructure:"last_updated"`
+        Tags           []*Tag                 `json:"tags" mapstructure:"tags"`
+        CustomFields   map[string]interface{} `json:"custom_fields" mapstructure:"custom_fields"`
+}