From f21d10aac8a16383356e354839c42ec4447312d0 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <flamingcow@gmail.com>
Date: Sun, 26 Jul 2020 18:47:38 +0000
Subject: [PATCH] Path calculation, network diagram

---
 main.go  | 40 ++++++++++++++++++++------------------
 unifi.go | 59 ++++++++++++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 74 insertions(+), 25 deletions(-)

diff --git a/main.go b/main.go
index 9fd2dc4..482305f 100644
--- a/main.go
+++ b/main.go
@@ -25,20 +25,18 @@ func main() {
 	}
 	log.Printf("Network interface: %s", link.Attrs().Name)
 
+	upstream, err := getUpstream(link.Attrs().Name)
+	if err != nil {
+		log.Fatal(err)
+	}
+	log.Printf("Connected to %s on %s", upstream.Name, upstream.Port)
+
 	api, err := findAPI(link)
 	if err != nil {
 		log.Fatal(err)
 	}
 	log.Printf("UniFi API URL: %s", api)
 
-	/*
-	upstream, err := getUpstream(link.Attrs().Name)
-	if err != nil {
-		log.Fatal(err)
-	}
-	log.Printf("Connected to %s [%s] on %s", upstream.device, upstream.addr, upstream.port)
-	*/
-
 	unifi, err := NewClient(api)
 	if err != nil {
 		log.Fatal(err)
@@ -54,18 +52,22 @@ func main() {
 		log.Fatal(err)
 	}
 
-	for _, dev := range devices {
-		log.Printf("%#v", dev)
+	upstream_dev := devices[upstream.Name]
+	if upstream_dev == nil {
+		log.Fatalf("Can't find my upstream in UniFi device list")
 	}
+
+	upstream_dev.PathHop([]*Device{}, map[string]*Device{})
+	upstream_dev.LogHop("", true)
 }
 
-type upstream struct {
-	device string
-	port string
-	addr net.Addr
+type Upstream struct {
+	Name string
+	Port string
+	MAC string
 }
 
-func getUpstream(intName string) (*upstream, error) {
+func getUpstream(intName string) (*Upstream, error) {
 	intr, err := net.InterfaceByName(intName)
 	if err != nil {
 		return nil, err
@@ -95,16 +97,16 @@ func getUpstream(intName string) (*upstream, error) {
 		return nil, err
 	}
 
-	ret := &upstream {
-		addr: addr,
+	ret := &Upstream {
+		MAC: strings.ToUpper(addr.String()),
 	}
 
 	for _, tlv := range lldpFrame.Optional {
 		switch tlv.Type {
 		case lldp.TLVTypePortDescription:
-			ret.port = strings.ToLower(string(tlv.Value))
+			ret.Port = strings.ToLower(string(tlv.Value))
 		case lldp.TLVTypeSystemName:
-			ret.device = string(tlv.Value)
+			ret.Name = string(tlv.Value)
 		}
 	}
 
diff --git a/unifi.go b/unifi.go
index f5ab5ea..fc69535 100644
--- a/unifi.go
+++ b/unifi.go
@@ -11,6 +11,7 @@ import (
 	"net/http/cookiejar"
 	"os"
 	"path"
+	"sort"
 	"strings"
 	"time"
 
@@ -28,7 +29,9 @@ type Device struct {
 	MAC string `json:"mac"`
 
 	// Constructed fields
-	Neighbors map[string]*Device
+	Neighbors map[string]*Device // all device peers
+	Beyond map[string]*Device // device peers that we reach through this device
+	Path []*Device
 
 	// Intermediate data; ignore
 	Ports []*Port `json:"port_table"`
@@ -157,24 +160,30 @@ func (c *Client) ListDevices() (map[string]*Device, error) {
 	ret := map[string]*Device{}
 
 	for _, dev := range device_resp.Data {
+		dev.Neighbors = map[string]*Device{}
+		dev.Beyond = map[string]*Device{}
+		dev.Path = []*Device{}
 		dev.MAC = strings.ToUpper(dev.MAC)
-		ret[dev.MAC] = dev
+
+		if ret[dev.Name] != nil {
+			return nil, fmt.Errorf("Duplicate UniFi device names: %s", dev.Name)
+		}
+
+		ret[dev.Name] = dev
 	}
 
 	// Second loop after the table is populated
 	for _, dev := range ret {
-		dev.Neighbors = map[string]*Device{}
-
 		for _, port := range dev.Ports {
 			for _, lldp := range port.LLDP {
 				lldp.MAC = strings.ToUpper(lldp.MAC)
 
-				neighbor := ret[lldp.MAC]
+				neighbor := ret[lldp.Name]
 				if neighbor == nil {
 					continue
 				}
 
-				dev.Neighbors[neighbor.MAC] = neighbor
+				dev.Neighbors[neighbor.Name] = neighbor
 			}
 		}
 		dev.Ports = nil
@@ -183,6 +192,44 @@ func (c *Client) ListDevices() (map[string]*Device, error) {
 	return ret, nil
 }
 
+func (d *Device) PathHop(path []*Device, visited map[string]*Device) {
+	newPath := append(path, d)
+
+	d.Path = newPath
+	visited[d.Name] = d
+
+	for _, neighbor := range d.Neighbors {
+		if visited[neighbor.Name] != nil {
+			// Cycle
+			continue
+		}
+
+		d.Beyond[neighbor.Name] = neighbor
+		neighbor.PathHop(newPath, visited)
+	}
+}
+
+func (d *Device) LogHop(prefix string, last bool) {
+	log.Printf("%s+- %s", prefix, d.Name)
+
+	names := []string{}
+	for _, next := range d.Beyond {
+		names = append(names, next.Name)
+	}
+	sort.Strings(names)
+
+	i := 0
+	for _, name := range names {
+		next := d.Beyond[name]
+		if last {
+			next.LogHop(fmt.Sprintf("%s   ", prefix), i == len(names) - 1)
+		} else {
+			next.LogHop(fmt.Sprintf("%s|  ", prefix), i == len(names) - 1)
+		}
+		i++
+	}
+}
+
 func loadCreds() (*Creds, error) {
 	fh, err := os.OpenFile(credPath(), os.O_RDONLY, 0)
 	if err != nil {