From 8be479e69ef56ce9ffefe1cf82bdb8a286ca50f4 Mon Sep 17 00:00:00 2001
From: davidnewhall2 <david@sleepers.pro>
Date: Thu, 25 Jun 2020 02:04:27 -0700
Subject: [PATCH] A bit of cleanup

---
 core/unifi/events.go | 34 +++++++++++++++++++++++++---------
 core/unifi/ids.go    | 40 ++++++++++++++++++++--------------------
 core/unifi/types.go  |  4 ++++
 3 files changed, 49 insertions(+), 29 deletions(-)

diff --git a/core/unifi/events.go b/core/unifi/events.go
index 8b7da4da..1f5f8458 100644
--- a/core/unifi/events.go
+++ b/core/unifi/events.go
@@ -69,9 +69,6 @@ func (u *Unifi) GetSiteEvents(site *Site, hours time.Duration) ([]*Event, error)
 	return event.Data, nil
 }
 
-// Events satisfied the sort.Interface.
-type events []*Event
-
 // Event describes a UniFi Event.
 // API Path: /api/s/default/stat/event.
 type Event struct {
@@ -143,11 +140,6 @@ type Event struct {
 // IPGeo is part of the UniFi Event data. Each event may have up to three of these.
 // One for source, one for dest and one for the USG location.
 type IPGeo struct {
-	GeoIP
-}
-
-// GeoIP is a struct in a struct to deal with weird UniFi output.
-type GeoIP struct {
 	Asn           int64   `json:"asn"`
 	Latitude      float64 `json:"latitude"`
 	Longitude     float64 `json:"longitude"`
@@ -158,6 +150,9 @@ type GeoIP struct {
 	Organization  string  `json:"organization"`
 }
 
+// Events satisfied the sort.Interface.
+type events []*Event
+
 // Len satisfies sort.Interface.
 func (e events) Len() int {
 	return len(e)
@@ -180,5 +175,26 @@ func (v *IPGeo) UnmarshalJSON(data []byte) error {
 		return nil // it's empty
 	}
 
-	return json.Unmarshal(data, &v.GeoIP)
+	g := struct {
+		Asn           int64   `json:"asn"`
+		Latitude      float64 `json:"latitude"`
+		Longitude     float64 `json:"longitude"`
+		City          string  `json:"city"`
+		ContinentCode string  `json:"continent_code"`
+		CountryCode   string  `json:"country_code"`
+		CountryName   string  `json:"country_name"`
+		Organization  string  `json:"organization"`
+	}{}
+
+	err := json.Unmarshal(data, &g)
+	v.Asn = g.Asn
+	v.Latitude = g.Latitude
+	v.Longitude = g.Longitude
+	v.City = g.City
+	v.ContinentCode = g.ContinentCode
+	v.CountryCode = g.CountryCode
+	v.CountryName = g.CountryName
+	v.Organization = g.Organization
+
+	return err
 }
diff --git a/core/unifi/ids.go b/core/unifi/ids.go
index 359d7744..0140e157 100644
--- a/core/unifi/ids.go
+++ b/core/unifi/ids.go
@@ -7,19 +7,16 @@ import (
 	"time"
 )
 
-type idsList []*IDS
-
 // IDS holds an Intrusion Prevention System Event.
 type IDS struct {
 	Archived              FlexBool  `json:"archived"`
-	DstIPCountry          FlexBool  `json:"dstipCountry"`
 	DestPort              int       `json:"dest_port,omitempty"`
 	SrcPort               int       `json:"src_port,omitempty"`
+	FlowID                int64     `json:"flow_id"`
 	InnerAlertRev         int64     `json:"inner_alert_rev"`
 	InnerAlertSeverity    int64     `json:"inner_alert_severity"`
 	InnerAlertGID         int64     `json:"inner_alert_gid"`
 	InnerAlertSignatureID int64     `json:"inner_alert_signature_id"`
-	FlowID                int64     `json:"flow_id"`
 	Time                  int64     `json:"time"`
 	Timestamp             int64     `json:"timestamp"`
 	Datetime              time.Time `json:"datetime"`
@@ -28,6 +25,7 @@ type IDS struct {
 	DestIP                string    `json:"dest_ip"`
 	DstMAC                string    `json:"dst_mac"`
 	DstIPASN              string    `json:"dstipASN"`
+	DstIPCountry          string    `json:"dstipCountry"`
 	EventType             string    `json:"event_type"`
 	Host                  string    `json:"host"`
 	ID                    string    `json:"_id"`
@@ -42,9 +40,9 @@ type IDS struct {
 	SiteName              string    `json:"-"`
 	SourceName            string    `json:"-"`
 	SrcIP                 string    `json:"src_ip"`
-	SrcMAC                string    `json:"src_mac"`
 	SrcIPASN              string    `json:"srcipASN"`
 	SrcIPCountry          string    `json:"srcipCountry"`
+	SrcMAC                string    `json:"src_mac"`
 	Subsystem             string    `json:"subsystem"`
 	UniqueAlertID         string    `json:"unique_alertid"`
 	USGIP                 string    `json:"usgip"`
@@ -55,21 +53,6 @@ type IDS struct {
 	USGIPGeo              IPGeo     `json:"usgipGeo"`
 }
 
-// Len satisfies sort.Interface.
-func (e idsList) Len() int {
-	return len(e)
-}
-
-// Swap satisfies sort.Interface.
-func (e idsList) Swap(i, j int) {
-	e[i], e[j] = e[j], e[i]
-}
-
-// Less satisfies sort.Interface. Sort our list by Datetime.
-func (e idsList) Less(i, j int) bool {
-	return e[i].Datetime.Before(e[j].Datetime)
-}
-
 // GetIDS returns Intrusion Detection Systems events for a list of Sites.
 // timeRange may have a length of 0, 1 or 2. The first time is Start, the second is End.
 // Events between start and end are returned. End defaults to time.Now().
@@ -150,3 +133,20 @@ func makeEventParams(timeRange ...time.Time) (string, error) {
 
 	return string(params), err
 }
+
+type idsList []*IDS
+
+// Len satisfies sort.Interface.
+func (e idsList) Len() int {
+	return len(e)
+}
+
+// Swap satisfies sort.Interface.
+func (e idsList) Swap(i, j int) {
+	e[i], e[j] = e[j], e[i]
+}
+
+// Less satisfies sort.Interface. Sort our list by Datetime.
+func (e idsList) Less(i, j int) bool {
+	return e[i].Datetime.Before(e[j].Datetime)
+}
diff --git a/core/unifi/types.go b/core/unifi/types.go
index cd126982..77c46b15 100644
--- a/core/unifi/types.go
+++ b/core/unifi/types.go
@@ -37,8 +37,12 @@ const (
 	APILoginPathNew string = "/api/auth/login"
 	// APIEventPathIDS returns Intrusion Detection/Prevention Systems Events
 	APIEventPathIDS string = "/api/s/%s/stat/ips/event"
+	// APIEventPathAlarms contains the site alarms.
+	APIEventPathAlarms string = "/api/s/%s/list/alarm"
 	// APIPrefixNew is the prefix added to the new API paths; except login. duh.
 	APIPrefixNew string = "/proxy/network"
+	// APIAnomaliesPath returns site anomalies.
+	APIAnomaliesPath string = "/api/s/%s/stat/anomalies"
 )
 
 // path returns the correct api path based on the new variable.