From 0cb7d3e51438d0f86cc3ec0551049c78ace3a375 Mon Sep 17 00:00:00 2001
From: Joshua Spence <josh@joshuaspence.com>
Date: Fri, 17 Sep 2021 11:16:56 +1000
Subject: [PATCH] Add `AdoptDevice` and `ForgetDevice` methods (#45)

* Add `AdoptDevice` and `ForgetDevice` methods

* Address comments
---
 fields/main.go            |  2 ++
 unifi/device.generated.go |  6 ++--
 unifi/device.go           | 61 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/fields/main.go b/fields/main.go
index 95b29f5..be9a262 100644
--- a/fields/main.go
+++ b/fields/main.go
@@ -148,6 +148,8 @@ func NewResource(version string, structName string, resourcePath string) *Resour
 		baseType.Fields[" Key"] = NewFieldInfo("Key", "key", "string", "", false, false, "")
 	case resource.StructName == "Device":
 		baseType.Fields[" MAC"] = NewFieldInfo("MAC", "mac", "string", "", true, false, "")
+		baseType.Fields["Adopted"] = NewFieldInfo("Adopted", "adopted", "bool", "", false, false, "")
+		baseType.Fields["State"] = NewFieldInfo("State", "state", "DeviceState", "", false, false, "")
 	case resource.StructName == "User":
 		baseType.Fields[" IP"] = NewFieldInfo("IP", "ip", "string", "non-generated field", true, false, "")
 		baseType.Fields[" DevIdOverride"] = NewFieldInfo("DevIdOverride", "dev_id_override", "int", "non-generated field", true, false, "")
diff --git a/unifi/device.generated.go b/unifi/device.generated.go
index 20df394..eef9203 100644
--- a/unifi/device.generated.go
+++ b/unifi/device.generated.go
@@ -28,6 +28,7 @@ type Device struct {
 
 	MAC string `json:"mac,omitempty"`
 
+	Adopted                     bool                      `json:"adopted"`
 	AtfEnabled                  bool                      `json:"atf_enabled,omitempty"`
 	BandsteeringMode            string                    `json:"bandsteering_mode,omitempty"` // off|equal|prefer_5g
 	BaresipAuthUser             string                    `json:"baresip_auth_user,omitempty"` // ^\+?[a-zA-Z0-9_.\-!~*'()]*
@@ -85,8 +86,9 @@ type Device struct {
 	RpsOverride                 DeviceRpsOverride         `json:"rps_override,omitempty"`
 	SnmpContact                 string                    `json:"snmp_contact,omitempty"`  // .{0,255}
 	SnmpLocation                string                    `json:"snmp_location,omitempty"` // .{0,255}
-	StpPriority                 string                    `json:"stp_priority,omitempty"`  // 0|4096|8192|12288|16384|20480|24576|28672|32768|36864|40960|45056|49152|53248|57344|61440
-	StpVersion                  string                    `json:"stp_version,omitempty"`   // stp|rstp|disabled
+	State                       DeviceState               `json:"state"`
+	StpPriority                 string                    `json:"stp_priority,omitempty"` // 0|4096|8192|12288|16384|20480|24576|28672|32768|36864|40960|45056|49152|53248|57344|61440
+	StpVersion                  string                    `json:"stp_version,omitempty"`  // stp|rstp|disabled
 	SwitchVLANEnabled           bool                      `json:"switch_vlan_enabled,omitempty"`
 	UbbPairName                 string                    `json:"ubb_pair_name,omitempty"` // .{1,128}
 	Volume                      int                       `json:"volume,omitempty"`        // [0-9]|[1-9][0-9]|100
diff --git a/unifi/device.go b/unifi/device.go
index 4cae075..53fbf0d 100644
--- a/unifi/device.go
+++ b/unifi/device.go
@@ -2,6 +2,24 @@ package unifi
 
 import (
 	"context"
+	"fmt"
+)
+
+type DeviceState int
+
+const (
+	DeviceStateUnknown          DeviceState = 0
+	DeviceStateConnected        DeviceState = 1
+	DeviceStatePending          DeviceState = 2
+	DeviceStateFirmwareMismatch DeviceState = 3
+	DeviceStateUpgrading        DeviceState = 4
+	DeviceStateProvisioning     DeviceState = 5
+	DeviceStateHeartbeatMissed  DeviceState = 6
+	DeviceStateAdopting         DeviceState = 7
+	DeviceStateDeleting         DeviceState = 8
+	DeviceStateInformError      DeviceState = 9
+	DeviceStateAdoptFailed      DeviceState = 10
+	DeviceStateIsolated         DeviceState = 11
 )
 
 func (c *Client) ListDevice(ctx context.Context, site string) ([]Device, error) {
@@ -39,3 +57,46 @@ func (c *Client) GetDevice(ctx context.Context, site, id string) (*Device, error
 
 	return nil, &NotFoundError{}
 }
+
+func (c *Client) AdoptDevice(ctx context.Context, site, mac string) error {
+	reqBody := struct {
+		Cmd string `json:"cmd"`
+		MAC string `json:"mac"`
+	}{
+		Cmd: "adopt",
+		MAC: mac,
+	}
+
+	var respBody struct {
+		Meta meta `json:"meta"`
+	}
+
+	err := c.do(ctx, "POST", fmt.Sprintf("s/%s/cmd/devmgr", site), reqBody, &respBody)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (c *Client) ForgetDevice(ctx context.Context, site, mac string) error {
+	reqBody := struct {
+		Cmd  string   `json:"cmd"`
+		MACs []string `json:"macs"`
+	}{
+		Cmd:  "delete-device",
+		MACs: []string{mac},
+	}
+
+	var respBody struct {
+		Meta meta     `json:"meta"`
+		Data []Device `json:"data"`
+	}
+
+	err := c.do(ctx, "POST", fmt.Sprintf("s/%s/cmd/sitemgr", site), reqBody, &respBody)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}