Merge pull request 'Added support for UnifiAPs' (#6) from unifi-controller into master

Reviewed-on: #6

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bytes"
 	"encoding/json"
 	"errors"
 	"flag"
@@ -8,7 +9,9 @@ import (
 	"io/ioutil"
 	"log"
 	"net/http"
+	"net/http/cookiejar"
 	"os"
+	"strconv"
 	"strings"
 	"time"
 )
@@ -17,11 +20,14 @@ import (
 
 const (
 	baseURL   = "https://unifi.freifunk-troisdorf.de/v2.1"
+	ucBaseURL = "https://unifi.freifunk-troisdorf.de:8443"
 	iso8601   = "2006-01-02T15:04:05-0700"
 )
 
 // flags
 var token = flag.String("token", "", "Defines the x-auth-token")
+var ucUser = flag.String("ucUser", "", "Defines the Unifi API User")
+var ucPass = flag.String("ucPass", "", "Defines the Unifi API Password")
 var version = "development"
 var delay time.Duration = 60 * time.Second
 
@@ -33,7 +39,12 @@ func main() {
 	if *token == "" {
 		log.Fatalln("Please specify an API token via the flag '-token'")
 	}
-
+	if *ucPass == "" {
+		log.Fatalln("Please specify an API Password via the flag '-ucPass'")
+	}
+	if *ucUser == "" {
+		log.Fatalln("Please specify an API User via the flag '-ucUser'")
+	}
 	// start API processing (runs in a loop)
 	go processAPIs()
 
@@ -41,14 +52,121 @@ func main() {
 	serveJSON()
 }
 
-func processAPIs() {
-	tick := time.Tick(delay)
-	for range tick {
+func lookupModels(model string) string {
+	switch model {
+	case "BZ2", "U2S48", "U2Sv2":
+		return "Unifi AP"
+	case "BZ2LR", "U2L48", "U2Lv2":
+		return "UniFi AP-LR"
+	case "U7E", "U7Ev2":
+		return "UniFi AP-AC"
+	case "U7HD", "U7SHD":
+		return "UniFi AP-HD"
+	case "UXSDM":
+		return "UniFi AP-BaseStationXG"
+	case "UCMSH":
+		return "AP-MeshXG"
+	case "U7MP":
+		return "AP-AC-Mesh-Pro"
+	case "U7LR":
+		return "UniFi AP-AC-LR"
+	case "U7LT":
+		return "UniFi AP-AC-Lite"
+	case "U7P":
+		return "UniFi AP-Pro"
+	case "U7MSH":
+		return "UniFi AP-AC-Mesh"
+	case "U7PG2":
+		return "UniFi AP-AC-Pro"
+	default:
+		return "Unifi Gerät"
+	}
+}
+
+func itob(i int) bool {
+	if i == 1 {
+		return true
+	}
+	return false
+}
+
+func processUcAPIs() []node {
+	var nodes []node
+	d, err := getDevices("ucDevices.json")
+	if err != nil {
+		log.Fatalln(err)
+	}
+	ucAPI := newAPI(*ucUser, *ucPass, ucBaseURL)
+	ucAPI.ucLogin()
+	sites, err := ucAPI.ucGetSites()
+	if err != nil {
+		panic(err)
+	}
+	devices, err := ucAPI.ucGetDevices(sites)
+	if err != nil {
+		panic(err)
+	}
+	for _, jsonDevice := range d.Devices {
+		var currentDevice ucDevice
+		var currentJSONDevice device
+		for _, device := range devices {
+			if strings.ToUpper(device.Mac) == strings.ToUpper(jsonDevice.MAC) {
+				currentDevice = device
+				currentJSONDevice = jsonDevice
+			}
+		}
+		load, err := strconv.ParseFloat(currentDevice.Sysstats.CPU, 64)
+		if err != nil {
+			panic(err)
+		}
+		mem, err := strconv.ParseFloat(currentDevice.Sysstats.Memory, 64)
+		if err != nil {
+			panic(err)
+		}
+
+		nodes = append(nodes, node{
+			Firstseen:      "0",
+			Lastseen:       time.Unix(int64(currentDevice.LastSeen), 0).Format(iso8601),
+			IsOnline:       itob(currentDevice.State),
+			IsGateway:      false,
+			Clients:        0,
+			ClientsWifi24:  0,
+			ClientsWifi5:   0,
+			ClientsOther:   0,
+			RootFSUsage:    0,
+			LoadAVG:        load / 100,
+			MemoryUsage:    mem / 100,
+			Uptime:         time.Now().Add(-1 * time.Second * time.Duration(currentDevice.Uptime)).Format(iso8601),
+			GatewayNexthop: currentJSONDevice.GatewayNexthop,
+			Gateway:        currentJSONDevice.Gateway,
+			Location:       currentJSONDevice.Location,
+			NodeID:         strings.ReplaceAll(currentDevice.Mac, ":", ""),
+			MAC:            currentDevice.Mac,
+			Adresses:       []string{currentDevice.IP},
+			Domain:         currentJSONDevice.Domain,
+			Hostname:       "[Unifi] " + currentDevice.Name,
+			Owner:          "Freifunk Rhein-Sieg",
+			Firmware: firmware{
+				Base:    "Ubiquiti - Stock",
+				Release: currentDevice.Version,
+			},
+			Autoupdater: autoupdater{
+				Enabled: false,
+				Branch:  "stable",
+			},
+			NProc: 1,
+			Model: lookupModels(currentDevice.Model),
+		})
+	}
+	return nodes
+}
+
+func processUNMSAPI() ([]node, []link) {
 	// Variables for runtime
 	var links []link
 	var nodes []node
 
-		d, err := getDevices()
+	d, err := getDevices("devices.json")
 	if err != nil {
 		log.Fatalln(err)
 	}
@@ -130,6 +248,19 @@ func processAPIs() {
 			Model: details.Identification.Model,
 		})
 	}
+	return nodes, links
+}
+
+func processAPIs() {
+	tick := time.Tick(delay)
+	for range tick {
+		var nodes []node
+		var links []link
+
+		unmsNodes, links := processUNMSAPI()
+		ucNodes := processUcAPIs()
+		nodes = append(nodes, unmsNodes...)
+		nodes = append(nodes, ucNodes...)
 
 		// assemble final struct
 		o := output{
@@ -150,9 +281,9 @@ func processAPIs() {
 	}
 }
 
-func getDevices() (devices, error) {
+func getDevices(file string) (devices, error) {
 	// get devices from JSON file
-	jsonFile, err := os.Open("devices.json")
+	jsonFile, err := os.Open(file)
 	if err != nil {
 		return devices{}, errors.New("can't open devices.json")
 	}
@@ -242,3 +373,88 @@ func serveJSON() {
 		log.Fatalln(err)
 	}
 }
+
+func httpClient() *http.Client {
+	jar, err := cookiejar.New(nil)
+	if err != nil {
+		log.Fatal(err)
+	}
+	client := &http.Client{Jar: jar}
+	return client
+}
+
+func newAPI(user string, pass string, baseURL string) ucAPIData {
+	return ucAPIData{
+		user:    user,
+		pass:    pass,
+		baseURL: baseURL,
+		client:  httpClient(),
+	}
+}
+
+func (u *ucAPIData) ucCallAPI(url string, method string, body *bytes.Buffer, output interface{}) error {
+	req, err := http.NewRequest(method, u.baseURL+url, body)
+	if err != nil {
+		return fmt.Errorf("can't set request %s", u.baseURL+url)
+	}
+
+	req.Header.Set("Content-Type", "application/json")
+	response, err := u.client.Do(req)
+	if err != nil {
+		return fmt.Errorf("can't login %s", u.baseURL+url)
+	}
+	defer response.Body.Close()
+	if response.StatusCode != 200 {
+		return fmt.Errorf("Login failed %s", u.baseURL+url)
+	}
+
+	data, err := ioutil.ReadAll(response.Body)
+	if err != nil {
+		return err
+	}
+
+	err = json.Unmarshal(data, &output)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (u *ucAPIData) ucLogin() error {
+	var loginData = []byte(`{"username":"` + u.user + `","password":"` + u.pass + `"}`)
+
+	url := "/api/login"
+	err := u.ucCallAPI(url, http.MethodPost, bytes.NewBuffer(loginData), nil)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (u *ucAPIData) ucGetSites() ([]ucSite, error) {
+	var d struct {
+		Data []ucSite `json:"data"`
+	}
+
+	url := "/api/self/sites"
+	err := u.ucCallAPI(url, http.MethodGet, bytes.NewBuffer([]byte{}), &d)
+	if err != nil {
+		return []ucSite{}, err
+	}
+	return d.Data, nil
+}
+
+func (u *ucAPIData) ucGetDevices(sites []ucSite) ([]ucDevice, error) {
+	var d struct {
+		Data []ucDevice `json:"data"`
+	}
+	var s []ucDevice
+
+	for _, site := range sites {
+		url := "/api/s/" + site.ID + "/stat/device"
+		u.ucCallAPI(url, http.MethodGet, bytes.NewBuffer([]byte{}), &d)
+		s = append(s, d.Data...)
+	}
+	return s, nil
+}

types.go

@@ -4,7 +4,9 @@ import (
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
+	"net/http"
 	"os"
+	"sync"
 	"time"
 )
 
@@ -142,3 +144,36 @@ type apiResponse struct {
 	Error      string `json:"error"`
 	Message    string `json:"message"`
 }
+
+type jar struct {
+	lk      sync.Mutex
+	cookies map[string][]*http.Cookie
+}
+
+type ucSite struct {
+	Name string `json:"desc"`
+	ID   string `json:"name"`
+}
+
+type ucDevice struct {
+	IP       string `json:"ip"`
+	Mac      string `json:"mac"`
+	Model    string `json:"model"`
+	Version  string `json:"version"`
+	Gateway  string `json:"gateway_mac"`
+	Name     string `json:"name"`
+	State    int    `json:"state"`
+	LastSeen int    `json:"last_seen"`
+	Uptime   int    `json:"uptime"`
+	Sysstats struct {
+		CPU    string `json:"cpu"`
+		Memory string `json:"mem"`
+	} `json:"system-stats"`
+}
+
+type ucAPIData struct {
+	user    string
+	pass    string
+	baseURL string
+	client  *http.Client
+}

ucDevices.json

@@ -0,0 +1,26 @@
+{
+    "devices":[
+      {
+        "name": "UAP-AC-M_S03",
+        "mac": "18:e8:29:56:6d:9e",
+        "gateway_nexthop": "18e8292f7de6",
+        "gateway": "a28cae6ff604",
+        "domain": "unifi",
+        "location": {
+          "longitude":7.148406208,
+          "latitude":50.817093402
+        }
+      },
+      {
+        "name": "Am-Krausacker-2",
+        "mac": "18:e8:29:a0:6f:23",
+        "gateway_nexthop": "18e8292f7de6",
+        "gateway": "a28cae6ff604",
+        "domain": "unifi",
+        "location": {
+          "longitude":7.148406208,
+          "latitude":50.817093402
+        }
+      }
+    ]
+  }