topo/graph/graph.go

package graph

import (
	"context"
	"fmt"
	"sync"

	"github.com/digitalocean/go-netbox/netbox/client"
	"github.com/digitalocean/go-netbox/netbox/client/dcim"
	"github.com/digitalocean/go-netbox/netbox/models"
	"github.com/golang/glog"

	confpb "code.hackerspace.pl/q3k/topo/proto/config"
)

type MachinePort struct {
	Machine  *Machine
	OtherEnd *SwitchPort
	Name     string
}

type SwitchPort struct {
	Switch   *Switch
	OtherEnd *MachinePort
	Name     string
}

type Machine struct {
	Name     string
	Complete bool

	Ports map[string]*MachinePort
}

type Switch struct {
	Name     string
	Complete bool

	Ports map[string]*SwitchPort
}

type Graph struct {
	Switches map[string]*Switch
	Machines map[string]*Machine
	Mu       sync.RWMutex
}

func New() *Graph {
	return &Graph{
		Switches: make(map[string]*Switch),
		Machines: make(map[string]*Machine),
	}
}

func (g *Graph) RemoveMachine(name string) {
	glog.Infof("Removed machine %q", name)
}

func (g *Graph) RemoveSwitch(name string) {
	glog.Infof("Removed switch %q", name)
}

func (g *Graph) LoadConfig(conf *confpb.Config) error {
	loadedMachines := make(map[string]bool)
	loadedSwitches := make(map[string]bool)

	// Add new machines and switches.
	for _, machinepb := range conf.Machine {
		if machinepb.Name == "" {
			return fmt.Errorf("empty machine name")
		}
		if loadedMachines[machinepb.Name] {
			return fmt.Errorf("duplicate machine name: %v", machinepb.Name)
		}
		machine, ok := g.Machines[machinepb.Name]
		if !ok {
			machine = &Machine{
				Name:  machinepb.Name,
				Ports: make(map[string]*MachinePort),
			}
			for _, portpb := range machinepb.ManagedPort {
				machine.Ports[portpb.Name] = &MachinePort{
					Name:    portpb.Name,
					Machine: machine,
				}
			}
			g.Machines[machinepb.Name] = machine
			glog.Infof("Added machine %q with %d managed ports", machine.Name, len(machine.Ports))
		}
		machine.Complete = false
		loadedMachines[machinepb.Name] = true
	}
	for _, switchpb := range conf.Switch {
		if switchpb.Name == "" {
			return fmt.Errorf("empty switch name")
		}
		if loadedSwitches[switchpb.Name] {
			return fmt.Errorf("duplicate switch name: %v", switchpb.Name)
		}
		if loadedMachines[switchpb.Name] {
			return fmt.Errorf("switch name collides with machine name: %v", switchpb.Name)
		}
		sw, ok := g.Switches[switchpb.Name]
		if !ok {
			sw = &Switch{
				Name:  switchpb.Name,
				Ports: make(map[string]*SwitchPort),
			}
			for _, portpb := range switchpb.ManagedPort {
				sw.Ports[portpb.Name] = &SwitchPort{
					Name:   portpb.Name,
					Switch: sw,
				}
			}
			g.Switches[switchpb.Name] = sw
			glog.Infof("Added switch %q with %d managed ports", sw.Name, len(sw.Ports))
		}
		sw.Complete = false
		loadedSwitches[switchpb.Name] = true
	}

	// Remove old machines and switches.
	removeMachines := make(map[string]bool)
	removeSwitches := make(map[string]bool)
	for name, _ := range g.Switches {
		if !loadedSwitches[name] {
			removeSwitches[name] = true
		}
	}
	for name, _ := range g.Machines {
		if !loadedMachines[name] {
			removeMachines[name] = true
		}
	}
	for name, _ := range removeMachines {
		g.RemoveMachine(name)
	}
	for name, _ := range removeSwitches {
		g.RemoveSwitch(name)
	}
	return nil

}

func (g *Graph) FeedFromNetbox(ctx context.Context, nb *client.NetBox) error {
	// Clear all connections first, because it's easier that way.
	for _, machine := range g.Machines {
		for _, port := range machine.Ports {
			port.OtherEnd = nil
		}
	}
	for _, sw := range g.Switches {
		for _, port := range sw.Ports {
			port.OtherEnd = nil
		}
	}

	// Load new connections.
	// Iterating over just machines should be fine if all connections are
	// guaranteed to be between machines and switches (which is the model for
	// now).
	for _, machine := range g.Machines {
		req := &dcim.DcimInterfaceConnectionsListParams{
			Device:  &machine.Name,
			Context: ctx,
		}
		res, err := nb.Dcim.DcimInterfaceConnectionsList(req, nil)
		if err != nil {
			return fmt.Errorf("while querying information about %q: %v", machine.Name, err)
		}
		for _, connection := range res.Payload.Results {
			ia := connection.InterfaceA
			ib := connection.InterfaceB
			if ia == nil || ib == nil {
				continue
			}

			// Find which way this thing actually connects.
			var thisSide, otherSide *models.PeerInterface
			if ia.Device.Name == machine.Name {
				thisSide = ia
				otherSide = ib
			} else if ib.Device.Name == machine.Name {
				thisSide = ib
				otherSide = ia
			} else {
				glog.Warning("Netbox connectivity for %q reported a link without it involced..?", machine.Name)
				continue
			}

			thisPort, ok := machine.Ports[*thisSide.Name]
			if !ok {
				continue
			}
			sw, ok := g.Switches[otherSide.Device.Name]
			if !ok {
				glog.Warningf("Machine %q port %q is managed but connected to unknown device %q", machine.Name, thisPort.Name, otherSide.Device.Name)
				continue
			}
			otherPort, ok := sw.Ports[*otherSide.Name]
			if !ok {
				glog.Warningf("Machine %q port %q is managed but connected to unmanaged port %q on %q", machine.Name, thisPort.Name, otherSide.Name, sw.Name)
				continue
			}

			// Connect the two together!
			thisPort.OtherEnd = otherPort
			otherPort.OtherEnd = thisPort
			glog.Infof("Connected: %s/%s <-> %s/%s", machine.Name, thisPort.Name, sw.Name, otherPort.Name)
		}
	}
	return nil
}