Add test for lookup.LookupByRoles

inverse-inc · Oct 1, 2024 · 65f8b7b · 65f8b7b
1 parent 97951a3
commit 65f8b7b
Show file tree

Hide file tree

Showing 3 changed files with 105 additions and 17 deletions.
diff --git a/go/cron/network_event.go b/go/cron/network_event.go
@@ -195,16 +195,23 @@ type NetworkTranslationInfo struct {
 }
 
 func (ne *NetworkEvent) GetSrcRole(ctx context.Context, db *sql.DB) (string, string) {
-	src := ne.SourceInventoryItem
-	if src == nil {
+	return ne.getRoleFromInventory(ctx, db, ne.SourceInventoryItem)
+}
+
+func (ne *NetworkEvent) GetDstRole(ctx context.Context, db *sql.DB) (string, string) {
+	return ne.getRoleFromInventory(ctx, db, ne.DestInventoryitem)
+}
+
+func (ne *NetworkEvent) getRoleFromInventory(ctx context.Context, db *sql.DB, item *InventoryItem) (string, string) {
+	if item == nil {
 		return "", ""
 	}
 
-	if len(src.ExternalIDS) == 0 {
+	if len(item.ExternalIDS) == 0 {
 		return "", ""
 	}
 
-	mac := src.ExternalIDS[0]
+	mac := item.ExternalIDS[0]
 	if mac == "" || mac == "00:00:00:00:00:00" {
 		return "", ""
 	}

diff --git a/go/cron/policy_lookup.go b/go/cron/policy_lookup.go
@@ -247,26 +247,55 @@ type PolicyLookup struct {
 }
 
 func (l PolicyLookup) Lookup(ctx context.Context, db *sql.DB, ne *NetworkEvent) *EnforcementInfo {
-	mac, role := ne.GetSrcRole(ctx, db)
-	if mac == "" || role == "" {
+	srcMac, srcRole := ne.GetSrcRole(ctx, db)
+	dstMac, dstRole := ne.GetDstRole(ctx, db)
+	if ei := l.LookupByMac(srcMac, ne); ei != nil {
+		return ei
+	}
+
+	if ei := l.LookupByMac(dstMac, ne); ei != nil {
+		return ei
+	}
+
+	if ei := l.LookupByRoles(srcRole, ne); ei != nil {
+		return ei
+	}
+
+	if ei := l.LookupByRoles(dstRole, ne); ei != nil {
+		return ei
+	}
+
+	if srcMac != "" {
+		return l.LookupImplict(ne)
+	}
+
+	return nil
+}
+
+func (l *PolicyLookup) LookupByRoles(role string, ne *NetworkEvent) *EnforcementInfo {
+	policies, ok := l.ByRoles[role]
+	if !ok {
 		return nil
 	}
 
-	if policies, ok := l.NodesPolicies[mac]; ok {
-		ei := matchEnforcementInfo(policies, ne)
-		if ei != nil {
-			return ei
-		}
+	if ei := matchEnforcementInfo(policies, ne); ei != nil {
+		return ei
 	}
 
-	if policies, ok := l.ByRoles[role]; ok {
-		ei := matchEnforcementInfo(policies, ne)
-		if ei != nil {
-			return ei
-		}
+	return nil
+}
+
+func (l *PolicyLookup) LookupByMac(mac string, ne *NetworkEvent) *EnforcementInfo {
+	policies, ok := l.NodesPolicies[mac]
+	if !ok {
+		return nil
 	}
 
-	return l.LookupImplict(ne)
+	if ei := matchEnforcementInfo(policies, ne); ei != nil {
+		return ei
+	}
+
+	return nil
 }
 
 func matchEnforcementInfo(policies []Policy, ne *NetworkEvent) *EnforcementInfo {

diff --git a/go/cron/policy_lookup_test.go b/go/cron/policy_lookup_test.go
@@ -562,6 +562,36 @@ const RolesPoliciesMapJSON = `
           "permit tcp any 10.15.1.0 0.0.0.255 eq 3389",
           "permit udp any 10.15.1.0 0.0.0.255 eq 3389"
         ]
+      },
+      {
+        "enforcement_info": [
+          {
+            "policy-revision": 3,
+            "verdict": "allow",
+            "dc-inventory-revision": 1725462233,
+            "rule-id": "28477cf7-234e-4751-8ced-542464017b1c/"
+          }
+        ],
+        "acls": [
+          "permit tcp any 10.15.1.0 0.0.0.255 eq 3389",
+          "permit udp any 10.15.1.0 0.0.0.255 eq 3389"
+        ]
+      },
+      {
+        "enforcement_info": [
+          {
+            "policy-revision": 66,
+            "verdict": "allow",
+            "dc-inventory-revision": 1727715416,
+            "rule-id": "d2cdcbd9-5acd-4021-ba96-fdecbbf77473/"
+          }
+        ],
+        "acls": [
+          "#permit tcp any host 00:50:56:9d:44:ca eq 222",
+          "#permit udp any host 00:50:56:9d:44:ca eq 222",
+          "#permit tcp any host 00:50:56:9d:44:ca eq 333",
+          "#permit udp any host 00:50:56:9d:44:ca eq 333"
+        ]
       }
     ]
   },
@@ -608,4 +638,26 @@ func TestPolicyLoad(t *testing.T) {
 	}
 
 	lookup.UpdateMatchers()
+	ne := NetworkEvent{
+		DestPort:   222,
+		SourceIp:   netip.AddrFrom4([4]byte{10, 0, 0, 1}),
+		DestIp:     netip.AddrFrom4([4]byte{10, 0, 0, 3}),
+		IpProtocol: IpProtocolUdp,
+		DestInventoryitem: &InventoryItem{
+			ExternalIDS: []string{"00:50:56:9d:44:ca"},
+		},
+	}
+
+	if diff := cmp.Diff(
+		lookup.LookupByRoles("IoT-Lighting", &ne),
+		&EnforcementInfo{
+			RuleID:              "d2cdcbd9-5acd-4021-ba96-fdecbbf77473/",
+			Verdict:             "allow",
+			PolicyRevision:      66,
+			DcInventoryRevision: 1727715416,
+		},
+	); diff != "" {
+		t.Fatalf("LookupByRoles does not match %s", diff)
+	}
+
 }