move string formatting to analyzer

Signed-off-by: Elazar Gershuni <[email protected]>
np-guard · Jun 19, 2024 · 6d566b3 · 6d566b3
1 parent 3c3f10a
commit 6d566b3
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 131 deletions.
diff --git a/pkg/connection/connection.go b/pkg/connection/connection.go
@@ -7,25 +7,10 @@ SPDX-License-Identifier: Apache-2.0
 package connection
 
 import (
-	"fmt"
-	"sort"
-	"strings"
-
-	"github.com/np-guard/models/pkg/ds"
 	"github.com/np-guard/models/pkg/netp"
 	"github.com/np-guard/models/pkg/netset"
 )
 
-const (
-	AllConnections = "All Connections"
-	NoConnections  = "No Connections"
-)
-const (
-	spaceString    = " "
-	protocolString = "protocol: "
-	commaString    = ", "
-)
-
 type Set = netset.TransportSet
 
 func NewTCPorUDP(protocol netp.ProtocolString, srcMinP, srcMaxP, dstMinP, dstMaxP int64) *Set {
@@ -41,23 +26,9 @@ func NewTCPSet() *Set {
 }
 
 func ICMPConnection(icmpType, icmpCode *int64) (*Set, error) {
-	icmp := netp.ICMP{}
-	if icmpType != nil {
-		t := int(*icmpType)
-		typeCode := netp.ICMPTypeCode{Type: t}
-		if netp.MaxCodes[t] == 0 {
-			var zero int64 = 0
-			icmpCode = &zero
-		}
-		if icmpCode != nil {
-			code := int(*icmpCode)
-			typeCode.Code = &code
-		}
-		err := netp.ValidateICMP(&typeCode)
-		if err != nil {
-			return nil, err
-		}
-		icmp = netp.ICMP{TypeCode: &typeCode}
+	icmp, err := netp.ICMPFromTypeAndCode(icmpType, icmpCode)
+	if err != nil {
+		return nil, err
 	}
 	return netset.NewICMPTransport(icmp), nil
 }
@@ -69,96 +40,3 @@ func All() *Set {
 func None() *Set {
 	return netset.AllOrNothingTransport(false, false)
 }
-
-func tcpudpProtocolString(p *netset.ProtocolSet, shortVersion bool) string {
-	var str string
-	switch {
-	case p.Size() == 2:
-		str = string(netp.ProtocolStringTCP + "," + netp.ProtocolStringUDP)
-	case p.Contains(netset.TCPCode):
-		str = string(netp.ProtocolStringTCP)
-	case p.Contains(netset.UDPCode):
-		str = string(netp.ProtocolStringUDP)
-	default:
-		return ""
-	}
-	if shortVersion {
-		return str
-	}
-	return protocolString + str
-}
-
-func getTCPUDPCubeStr(cube ds.Triple[*netset.ProtocolSet, *netset.PortSet, *netset.PortSet], shortVersion bool) string {
-	var ports []string
-	if !cube.S2.Equal(netset.AllPorts()) {
-		ports = append(ports, "src-ports: "+cube.S2.String())
-	}
-	if !cube.S3.Equal(netset.AllPorts()) {
-		ports = append(ports, "dst-ports: "+cube.S3.String())
-	}
-	res := tcpudpProtocolString(cube.S1, shortVersion)
-	if len(ports) > 0 {
-		res += spaceString
-	}
-	return res + strings.Join(ports, spaceString)
-}
-
-func getICMPCubeStr(cube netp.ICMP) string {
-	tc := cube.ICMPTypeCode()
-	if tc == nil {
-		return ""
-	}
-	if tc.Code == nil {
-		if netp.MaxCodes[tc.Type] == 0 {
-			return fmt.Sprintf("icmp-type: %d icmp-code: 0", tc.Type)
-		}
-		return fmt.Sprintf("icmp-type: %d", tc.Type)
-	}
-	return fmt.Sprintf("icmp-type: %d icmp-code: %d", tc.Type, *tc.Code)
-}
-
-// tCPUDPString returns a string representation of a TransportSet object
-func tCPUDPString(c *netset.TCPUDPSet, shortVersion bool) string {
-	cubes := c.Partitions()
-	var resStrings = make([]string, len(cubes))
-	for i, cube := range cubes {
-		resStrings[i] = getTCPUDPCubeStr(cube, shortVersion)
-	}
-	return strings.Join(resStrings, commaString)
-}
-
-// iCMPString returns a string representation of an ICMPSet object
-func iCMPString(c *netset.ICMPSet, shortVersion bool) string {
-	if c.IsEmpty() {
-		return ""
-	}
-	cubes := c.Partitions()
-	var resStrings = make([]string, len(cubes))
-	for i, cube := range cubes {
-		resStrings[i] = getICMPCubeStr(cube)
-	}
-	sort.Strings(resStrings)
-	str := "ICMP"
-	if !shortVersion {
-		str = protocolString + str
-	}
-	last := strings.Join(resStrings, commaString)
-	if last != "" {
-		str += spaceString + last
-	}
-	return str
-}
-
-func Stringify(c *Set, shortVersion bool) string {
-	if c.IsEmpty() {
-		return NoConnections
-	} else if c.IsAll() {
-		return AllConnections
-	}
-	tcpString := tCPUDPString(c.TCPUDPSet(), shortVersion)
-	icmpString := iCMPString(c.ICMPSet(), shortVersion)
-	if tcpString != "" && icmpString != "" {
-		return fmt.Sprintf("%s; %s", icmpString, tcpString)
-	}
-	return icmpString + tcpString
-}
diff --git a/pkg/netp/icmp.go b/pkg/netp/icmp.go
@@ -31,6 +31,28 @@ func NewICMP(typeCode *ICMPTypeCode) (ICMP, error) {
 	return ICMP{TypeCode: typeCode}, nil
 }
 
+func ICMPFromTypeAndCode(icmpType, icmpCode *int64) (ICMP, error) {
+	icmp := ICMP{}
+	if icmpType != nil {
+		t := int(*icmpType)
+		typeCode := ICMPTypeCode{Type: t}
+		if HasSingleCode(t) {
+			var zero int64 = 0
+			icmpCode = &zero
+		}
+		if icmpCode != nil {
+			code := int(*icmpCode)
+			typeCode.Code = &code
+		}
+		err := ValidateICMP(&typeCode)
+		if err != nil {
+			return icmp, err
+		}
+		icmp = ICMP{TypeCode: &typeCode}
+	}
+	return icmp, nil
+}
+
 func (t ICMP) ICMPTypeCode() *ICMPTypeCode {
 	if t.TypeCode == nil {
 		return nil
@@ -99,7 +121,7 @@ func inverseICMPType(t int) int {
 	return undefinedICMP
 }
 
-var MaxCodes = map[int]int{
+var maxCodes = map[int]int{
 	EchoReply:              0,
 	DestinationUnreachable: 5,
 	SourceQuench:           0,
@@ -113,11 +135,23 @@ var MaxCodes = map[int]int{
 	InformationReply:       0,
 }
 
+func MaxCode(t int) int {
+	return maxCodes[t]
+}
+
+func Types() []int {
+	var types []int
+	for t := range maxCodes {
+		types = append(types, t)
+	}
+	return types
+}
+
 func ValidateICMP(typeCode *ICMPTypeCode) error {
 	if typeCode == nil {
 		return nil
 	}
-	maxCode, ok := MaxCodes[typeCode.Type]
+	maxCode, ok := maxCodes[typeCode.Type]
 	if !ok {
 		return fmt.Errorf("invalid ICMP type %v", typeCode.Type)
 	}
@@ -127,6 +161,10 @@ func ValidateICMP(typeCode *ICMPTypeCode) error {
 	return nil
 }
 
+func HasSingleCode(t int) bool {
+	return maxCodes[t] == 0
+}
+
 func (t ICMP) ProtocolString() ProtocolString {
 	return ProtocolStringICMP
 }
diff --git a/pkg/netset/icmpset.go b/pkg/netset/icmpset.go
@@ -125,12 +125,12 @@ func (s *ICMPSet) IntervalSet() *interval.CanonicalSet {
 
 func (s *ICMPSet) collect(old int) []netp.ICMP {
 	var res []netp.ICMP
-	for code := 0; code <= netp.MaxCodes[old]; code++ {
+	for code := 0; code <= netp.MaxCode(old); code++ {
 		if s.Contains(mapToNew(old, code)) {
 			res = append(res, netp.ICMP{TypeCode: &netp.ICMPTypeCode{Type: old, Code: &code}})
 		}
 	}
-	if len(res) == netp.MaxCodes[old]+1 {
+	if len(res) == netp.MaxCode(old)+1 {
 		res = []netp.ICMP{{TypeCode: &netp.ICMPTypeCode{Type: old, Code: nil}}}
 	}
 	return res
@@ -142,7 +142,7 @@ func (s *ICMPSet) Partitions() []netp.ICMP {
 		return []netp.ICMP{{TypeCode: nil}}
 	}
 	var res []netp.ICMP
-	for t := range netp.MaxCodes {
+	for t := range netp.Types() {
 		res = append(res, s.collect(t)...)
 	}
 	return res

diff --git a/pkg/netset/ipblock.go b/pkg/netset/ipblock.go
@@ -167,7 +167,7 @@ func DisjointIPBlocks(set1, set2 []*IPBlock) []*IPBlock {
 	for i, ipb := range set2 {
 		ipbList[len(set1)+i] = ipb.Copy()
 	}
-	// sort ipbList by ip_count per ipblock
+	// sort ipbList by ip_count per netset
 	sort.Slice(ipbList, func(i, j int) bool {
 		return ipbList[i].ipCount() < ipbList[j].ipCount()
 	})