feat: Add new helper funcs spec

sandialabs · Sep 3, 2024 · 03ce90b · 03ce90b
1 parent 0264514
commit 03ce90b
Show file tree

Hide file tree

Showing 7 changed files with 214 additions and 0 deletions.
diff --git a/src/go/types/interfaces/topology.go b/src/go/types/interfaces/topology.go
@@ -10,6 +10,7 @@ type TopologySpec interface {
 	FindNodeByName(string) NodeSpec
 	FindNodesWithLabels(...string) []NodeSpec
 	FindDelayedNodes() []NodeSpec
+	FindNodesWithVLAN(string) []NodeSpec
 
 	AddNode(string, string) NodeSpec
 	RemoveNode(string)
@@ -35,11 +36,55 @@ type NodeSpec interface {
 	External() bool
 
 	SetInjections([]NodeInjection)
+	SetType(string)
 
+	AddAnnotation(string, interface{})
+	AddTimerDelay(delay string)
+	AddUserDelay(delay bool)
+	AddC2Delay(hostname string, useuuid bool)
 	AddLabel(string, string)
 	AddHardware(string, int, int) NodeHardware
 	AddNetworkInterface(string, string, string) NodeNetworkInterface
 	AddNetworkRoute(string, string, int)
+	// Add NAT section to Node Network
+	//
+	// Example code:
+	//
+	//   newNat := map[string][]string{
+	//   	"eth0": {"eth1"},
+	//   }
+	//   nats := []map[string][]string{}
+	//   nats = append(nats, newNat)
+	//   node.AddNetworkNAT(nats)
+	//
+	// Example result:
+	//
+	//   nat:
+	//     - in:
+	//       - eth1
+	//       out: eth0
+	AddNetworkNAT(nats []map[string][]string)
+	// Add OSPF section to Node Network
+	//
+	// Example code:
+	//
+	//   area := map[int][]string{
+	//   	0: {"75.75.0.0/16"},
+	//   }
+	//   node.AddNetworkOSPF("75.75.75.5", 60, 10, 5, area)
+	//
+	// Example result:
+	//
+	//   ospf:
+	//     router_id: 75.75.75.5
+	//     dead_interval: 60
+	//     hello_interval: 10
+	//     retransmission_interval: 5
+	//     areas:
+	//       - area_id: 0
+	//         area_networks:
+	//       	 - network: 75.75.0.0/16
+	AddNetworkOSPF(routerID string, dead, hello, retrans int, areas map[int][]string)
 	AddInject(string, string, string, string)
 
 	SetAdvanced(map[string]string)
@@ -94,6 +139,8 @@ type NodeNetwork interface {
 	AddRuleset(NodeNetworkRuleset)
 
 	InterfaceAddress(string) string
+	InterfaceVLAN(string) string
+	InterfaceMask(string) int
 }
 
 type NodeNetworkInterface interface {

diff --git a/src/go/types/version/v0/network.go b/src/go/types/version/v0/network.go
@@ -77,6 +77,26 @@ func (this *Network) InterfaceAddress(name string) string {
 	return ""
 }
 
+func (this *Network) InterfaceVLAN(vlan string) string {
+	for _, iface := range this.InterfacesF {
+		if iface.VLAN() == vlan {
+			return iface.NameF
+		}
+	}
+
+	return ""
+}
+
+func (this *Network) InterfaceMask(name string) int {
+	for _, iface := range this.InterfacesF {
+		if strings.EqualFold(iface.NameF, name) {
+			return iface.MaskF
+		}
+	}
+
+	return 0
+}
+
 type Interface struct {
 	NameF       string   `json:"name" yaml:"name" structs:"name" mapstructure:"name"`
 	TypeF       string   `json:"type" yaml:"type" structs:"type" mapstructure:"type"`

diff --git a/src/go/types/version/v0/node.go b/src/go/types/version/v0/node.go
@@ -84,6 +84,24 @@ func (this *Node) SetInjections(injections []ifaces.NodeInjection) {
 	this.InjectionsF = injects
 }
 
+func (this *Node) SetType(t string) {
+	this.TypeF = t
+}
+
+func (this *Node) AddAnnotation(k string, i interface{}) {
+	if this.AnnotationsF == nil {
+		this.AnnotationsF = make(map[string]interface{})
+	}
+
+	this.AnnotationsF[k] = i
+}
+
+func (this *Node) AddTimerDelay(string) {}
+
+func (this *Node) AddUserDelay(bool) {}
+
+func (this *Node) AddC2Delay(string, bool) {}
+
 func (this *Node) AddLabel(k, v string) {
 	this.LabelsF[k] = v
 }
@@ -130,6 +148,26 @@ func (this *Node) AddNetworkRoute(dest, next string, cost int) {
 	this.NetworkF.RoutesF = append(this.NetworkF.RoutesF, r)
 }
 
+func (this *Node) AddNetworkNAT([]map[string][]string) {}
+
+func (this *Node) AddNetworkOSPF(routerID string, dead, hello, retrans int, areas map[int][]string) {
+	this.NetworkF.OSPFF = new(OSPF)
+	this.NetworkF.OSPFF.RouterIDF = routerID
+	this.NetworkF.OSPFF.DeadIntervalF = &dead
+	this.NetworkF.OSPFF.HelloIntervalF = &hello
+	this.NetworkF.OSPFF.RetransmissionIntervalF = &retrans
+
+	for id, networks := range areas {
+		area := new(Area)
+		area.AreaIDF = &id
+		for _, net := range networks {
+			areaNetwork := AreaNetwork{NetworkF: net}
+			area.AreaNetworksF = append(area.AreaNetworksF, areaNetwork)
+		}
+		this.NetworkF.OSPFF.AreasF = append(this.NetworkF.OSPFF.AreasF, *area)
+	}
+}
+
 func (this *Node) AddInject(src, dst, perms, desc string) {
 	this.InjectionsF = append(this.InjectionsF, &Injection{
 		SrcF:         src,

diff --git a/src/go/types/version/v0/topology.go b/src/go/types/version/v0/topology.go
@@ -94,6 +94,21 @@ func (TopologySpec) FindDelayedNodes() []ifaces.NodeSpec {
 	return nil
 }
 
+func (this TopologySpec) FindNodesWithVLAN(vlan string) []ifaces.NodeSpec {
+	var nodes []ifaces.NodeSpec
+
+	for _, n := range this.NodesF {
+		for _, i := range n.NetworkF.InterfacesF {
+			if i.VLAN() == vlan {
+				nodes = append(nodes, n)
+				break
+			}
+		}
+	}
+
+	return nodes
+}
+
 func (this *TopologySpec) AddNode(typ, hostname string) ifaces.NodeSpec {
 	n := &Node{
 		TypeF: typ,

diff --git a/src/go/types/version/v1/network.go b/src/go/types/version/v1/network.go
@@ -110,6 +110,26 @@ func (this *Network) InterfaceAddress(name string) string {
 	return ""
 }
 
+func (this *Network) InterfaceVLAN(vlan string) string {
+	for _, iface := range this.InterfacesF {
+		if iface.VLAN() == vlan {
+			return iface.NameF
+		}
+	}
+
+	return ""
+}
+
+func (this *Network) InterfaceMask(name string) int {
+	for _, iface := range this.InterfacesF {
+		if strings.EqualFold(iface.NameF, name) {
+			return iface.MaskF
+		}
+	}
+
+	return 0
+}
+
 type Interface struct {
 	NameF       string   `json:"name" yaml:"name" structs:"name" mapstructure:"name"`
 	TypeF       string   `json:"type" yaml:"type" structs:"type" mapstructure:"type"`

diff --git a/src/go/types/version/v1/node.go b/src/go/types/version/v1/node.go
@@ -102,6 +102,34 @@ func (this *Node) SetInjections(injections []ifaces.NodeInjection) {
 	this.InjectionsF = injects
 }
 
+func (this *Node) SetType(t string) {
+	this.TypeF = t
+}
+
+func (this *Node) AddAnnotation(k string, i interface{}) {
+	if this.AnnotationsF == nil {
+		this.AnnotationsF = make(map[string]interface{})
+	}
+
+	this.AnnotationsF[k] = i
+}
+
+func (this *Node) AddTimerDelay(delay string) {
+	this.DelayF.TimerF = delay
+}
+
+func (this *Node) AddUserDelay(delay bool) {
+	this.DelayF.UserF = delay
+}
+
+func (this *Node) AddC2Delay(hostname string, useuuid bool) {
+	this.DelayF = new(Delay)
+	d := new(C2Delay)
+	d.HostnameF = hostname
+	d.UseUUIDF = useuuid
+	this.DelayF.C2F = append(this.DelayF.C2F, *d)
+}
+
 func (this *Node) AddLabel(k, v string) {
 	if this.LabelsF == nil {
 		this.LabelsF = make(map[string]string)
@@ -152,6 +180,37 @@ func (this *Node) AddNetworkRoute(dest, next string, cost int) {
 	this.NetworkF.RoutesF = append(this.NetworkF.RoutesF, r)
 }
 
+func (this *Node) AddNetworkNAT(nats []map[string][]string) {
+	for _, nat := range nats {
+		n := new(NAT)
+		for out, ins := range nat {
+			n.OutF = out
+			for _, in := range ins {
+				n.InF = append(n.InF, in)
+			}
+		}
+		this.NetworkF.NATF = append(this.NetworkF.NATF, *n)
+	}
+}
+
+func (this *Node) AddNetworkOSPF(routerID string, dead, hello, retrans int, areas map[int][]string) {
+	this.NetworkF.OSPFF = new(OSPF)
+	this.NetworkF.OSPFF.RouterIDF = routerID
+	this.NetworkF.OSPFF.DeadIntervalF = &dead
+	this.NetworkF.OSPFF.HelloIntervalF = &hello
+	this.NetworkF.OSPFF.RetransmissionIntervalF = &retrans
+
+	for id, networks := range areas {
+		area := new(Area)
+		area.AreaIDF = &id
+		for _, net := range networks {
+			areaNetwork := AreaNetwork{NetworkF: net}
+			area.AreaNetworksF = append(area.AreaNetworksF, areaNetwork)
+		}
+		this.NetworkF.OSPFF.AreasF = append(this.NetworkF.OSPFF.AreasF, *area)
+	}
+}
+
 func (this *Node) AddInject(src, dst, perms, desc string) {
 	if _, ok := this.LabelsF["disable-injects"]; ok {
 		return

diff --git a/src/go/types/version/v1/topology.go b/src/go/types/version/v1/topology.go
@@ -105,6 +105,21 @@ func (this TopologySpec) FindDelayedNodes() []ifaces.NodeSpec {
 	return nodes
 }
 
+func (this TopologySpec) FindNodesWithVLAN(vlan string) []ifaces.NodeSpec {
+	var nodes []ifaces.NodeSpec
+
+	for _, n := range this.NodesF {
+		for _, i := range n.NetworkF.InterfacesF {
+			if i.VLAN() == vlan {
+				nodes = append(nodes, n)
+				break
+			}
+		}
+	}
+
+	return nodes
+}
+
 func (this *TopologySpec) AddNode(typ, hostname string) ifaces.NodeSpec {
 	n := &Node{
 		TypeF: typ,