From 4ec3c68be4da5bbe3229b147091554ed7fadc44b Mon Sep 17 00:00:00 2001 From: James Jeffries Date: Wed, 15 Oct 2025 16:31:56 +0100 Subject: [PATCH 1/2] handle policies as a string --- agent/configmgr/fleet/from_rpc.go | 17 ++ agent/configmgr/fleet/from_rpc_test.go | 229 +++++++++++++++++++++++++ 2 files changed, 246 insertions(+) diff --git a/agent/configmgr/fleet/from_rpc.go b/agent/configmgr/fleet/from_rpc.go index 1fa5785..0c62ab6 100644 --- a/agent/configmgr/fleet/from_rpc.go +++ b/agent/configmgr/fleet/from_rpc.go @@ -6,6 +6,8 @@ import ( "log/slog" "os" + "gopkg.in/yaml.v3" + "github.com/netboxlabs/orb-agent/agent/backend" "github.com/netboxlabs/orb-agent/agent/config" "github.com/netboxlabs/orb-agent/agent/configmgr/fleet/messages" @@ -159,6 +161,21 @@ func (messaging *Messaging) handleAgentPolicies(rpc []messages.AgentPolicyRPCPay for _, payload := range rpc { if payload.Action != "sanitize" { + // If the policy data is in YAML format and is a string, unmarshal it to a structured object + if payload.Format == "yaml" { + if dataStr, ok := payload.Data.(string); ok && dataStr != "" { + var structuredData map[string]any + if err := yaml.Unmarshal([]byte(dataStr), &structuredData); err != nil { + messaging.logger.Warn("failed to unmarshal YAML policy data", + "policy_id", payload.ID, + "policy_name", payload.Name, + "error", err) + // Continue with original data - let the backend handle the error + } else { + payload.Data = structuredData + } + } + } messaging.policyManager.ManagePolicy(config.PolicyPayload(payload)) } } diff --git a/agent/configmgr/fleet/from_rpc_test.go b/agent/configmgr/fleet/from_rpc_test.go index f508d44..762b224 100644 --- a/agent/configmgr/fleet/from_rpc_test.go +++ b/agent/configmgr/fleet/from_rpc_test.go @@ -1363,3 +1363,232 @@ func TestMessageHandlers_DispatchToHandlers_MalformedAgentStopPayload(t *testing // Assert assert.Error(t, err) } + +// Test handleAgentPolicies with YAML string data +func TestMessageHandlers_handleAgentPolicies_YAMLStringData(t *testing.T) { + // Arrange + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + mockPMgr := &mockPolicyManager{} + resetChan := make(chan struct{}, 1) + groupManager := newGroupManager() + handlers := NewMessaging(logger, mockPMgr, resetChan, &groupManager) + + // Setup mock expectations - verify that the data is converted to a map + mockPMgr.On("ManagePolicy", mock.MatchedBy(func(p config.PolicyPayload) bool { + if p.ID != "policy1" || p.Action != "manage" { + return false + } + // Verify that Data is now a map, not a string + dataMap, ok := p.Data.(map[string]any) + if !ok { + return false + } + // Verify the structure of the unmarshaled YAML + scope, ok := dataMap["scope"].(map[string]any) + if !ok { + return false + } + targets, ok := scope["targets"].([]any) + return ok && len(targets) == 1 && targets[0] == "192.168.12.190/32" + })).Return() + + // Create policy payload with YAML string data + yamlString := "scope:\n targets: [192.168.12.190/32]\n" + policies := []messages.AgentPolicyRPCPayload{ + { + Action: "manage", + ID: "policy1", + DatasetID: "dataset1", + AgentGroupID: "group1", + Name: "Network Discovery Policy", + Backend: "network-discovery", + Format: "yaml", + Version: 1, + Data: yamlString, + }, + } + + // Act + handlers.handleAgentPolicies(policies, false) + + // Assert + mockPMgr.AssertExpectations(t) +} + +// Test handleAgentPolicies with structured data (backwards compatibility) +func TestMessageHandlers_handleAgentPolicies_StructuredData(t *testing.T) { + // Arrange + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + mockPMgr := &mockPolicyManager{} + resetChan := make(chan struct{}, 1) + groupManager := newGroupManager() + handlers := NewMessaging(logger, mockPMgr, resetChan, &groupManager) + + // Setup mock expectations - verify that structured data is passed through unchanged + mockPMgr.On("ManagePolicy", mock.MatchedBy(func(p config.PolicyPayload) bool { + if p.ID != "policy2" || p.Action != "manage" { + return false + } + // Verify that Data remains as a map + dataMap, ok := p.Data.(map[string]any) + if !ok { + return false + } + scope, ok := dataMap["scope"].(map[string]any) + if !ok { + return false + } + targets, ok := scope["targets"].([]any) + return ok && len(targets) == 1 && targets[0] == "192.168.1.1/32" + })).Return() + + // Create policy payload with already structured data + structuredData := map[string]any{ + "scope": map[string]any{ + "targets": []any{"192.168.1.1/32"}, + }, + } + policies := []messages.AgentPolicyRPCPayload{ + { + Action: "manage", + ID: "policy2", + DatasetID: "dataset2", + AgentGroupID: "group1", + Name: "Device Discovery Policy", + Backend: "device-discovery", + Format: "yaml", + Version: 1, + Data: structuredData, + }, + } + + // Act + handlers.handleAgentPolicies(policies, false) + + // Assert + mockPMgr.AssertExpectations(t) +} + +// Test handleAgentPolicies with empty YAML string +func TestMessageHandlers_handleAgentPolicies_EmptyYAMLString(t *testing.T) { + // Arrange + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + mockPMgr := &mockPolicyManager{} + resetChan := make(chan struct{}, 1) + groupManager := newGroupManager() + handlers := NewMessaging(logger, mockPMgr, resetChan, &groupManager) + + // Setup mock expectations - empty string should be passed through + mockPMgr.On("ManagePolicy", mock.MatchedBy(func(p config.PolicyPayload) bool { + if p.ID != "policy3" { + return false + } + // Empty string should remain as empty string + _, ok := p.Data.(string) + return ok + })).Return() + + // Create policy payload with empty YAML string + policies := []messages.AgentPolicyRPCPayload{ + { + Action: "manage", + ID: "policy3", + DatasetID: "dataset3", + AgentGroupID: "group1", + Name: "Empty Policy", + Backend: "worker", + Format: "yaml", + Version: 1, + Data: "", + }, + } + + // Act + handlers.handleAgentPolicies(policies, false) + + // Assert + mockPMgr.AssertExpectations(t) +} + +// Test handleAgentPolicies with invalid YAML string +func TestMessageHandlers_handleAgentPolicies_InvalidYAMLString(t *testing.T) { + // Arrange + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + mockPMgr := &mockPolicyManager{} + resetChan := make(chan struct{}, 1) + groupManager := newGroupManager() + handlers := NewMessaging(logger, mockPMgr, resetChan, &groupManager) + + // Setup mock expectations - invalid YAML should be passed through as-is + mockPMgr.On("ManagePolicy", mock.MatchedBy(func(p config.PolicyPayload) bool { + if p.ID != "policy4" { + return false + } + // Invalid YAML should remain as string (unmarshal fails, original data used) + _, ok := p.Data.(string) + return ok + })).Return() + + // Create policy payload with invalid YAML string + invalidYAML := "this is not valid YAML: [[[{" + policies := []messages.AgentPolicyRPCPayload{ + { + Action: "manage", + ID: "policy4", + DatasetID: "dataset4", + AgentGroupID: "group1", + Name: "Invalid YAML Policy", + Backend: "pktvisor", + Format: "yaml", + Version: 1, + Data: invalidYAML, + }, + } + + // Act + handlers.handleAgentPolicies(policies, false) + + // Assert + mockPMgr.AssertExpectations(t) +} + +// Test handleAgentPolicies with non-yaml format +func TestMessageHandlers_handleAgentPolicies_NonYAMLFormat(t *testing.T) { + // Arrange + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + mockPMgr := &mockPolicyManager{} + resetChan := make(chan struct{}, 1) + groupManager := newGroupManager() + handlers := NewMessaging(logger, mockPMgr, resetChan, &groupManager) + + // Setup mock expectations - non-yaml format should pass data through unchanged + mockPMgr.On("ManagePolicy", mock.MatchedBy(func(p config.PolicyPayload) bool { + if p.ID != "policy5" { + return false + } + // Data should remain as string since format is not yaml + _, ok := p.Data.(string) + return ok + })).Return() + + // Create policy payload with non-yaml format + policies := []messages.AgentPolicyRPCPayload{ + { + Action: "manage", + ID: "policy5", + DatasetID: "dataset5", + AgentGroupID: "group1", + Name: "JSON Policy", + Backend: "pktvisor", + Format: "json", + Version: 1, + Data: "{\"key\": \"value\"}", + }, + } + + // Act + handlers.handleAgentPolicies(policies, false) + + // Assert + mockPMgr.AssertExpectations(t) +} From ebffc75f9d6bec6d95366157b5077e2d65c2ae35 Mon Sep 17 00:00:00 2001 From: James Jeffries Date: Wed, 15 Oct 2025 16:51:42 +0100 Subject: [PATCH 2/2] assume polciy is yaml if not specified --- agent/configmgr/fleet/from_rpc.go | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/agent/configmgr/fleet/from_rpc.go b/agent/configmgr/fleet/from_rpc.go index 0c62ab6..0b2987e 100644 --- a/agent/configmgr/fleet/from_rpc.go +++ b/agent/configmgr/fleet/from_rpc.go @@ -161,19 +161,22 @@ func (messaging *Messaging) handleAgentPolicies(rpc []messages.AgentPolicyRPCPay for _, payload := range rpc { if payload.Action != "sanitize" { - // If the policy data is in YAML format and is a string, unmarshal it to a structured object - if payload.Format == "yaml" { - if dataStr, ok := payload.Data.(string); ok && dataStr != "" { - var structuredData map[string]any - if err := yaml.Unmarshal([]byte(dataStr), &structuredData); err != nil { + // If the policy data is a string and Format is "yaml" (or empty), try to unmarshal it as YAML + // This handles cases where Format="yaml" or where the backend sends YAML without setting Format + if dataStr, ok := payload.Data.(string); ok && dataStr != "" && (payload.Format == "yaml" || payload.Format == "") { + var structuredData map[string]any + if err := yaml.Unmarshal([]byte(dataStr), &structuredData); err != nil { + // If unmarshaling fails, log a warning only if Format was explicitly set to yaml + if payload.Format == "yaml" { messaging.logger.Warn("failed to unmarshal YAML policy data", "policy_id", payload.ID, "policy_name", payload.Name, "error", err) - // Continue with original data - let the backend handle the error - } else { - payload.Data = structuredData } + // Continue with original string data - let the backend handle it + } else { + // Successfully unmarshaled - use the structured data + payload.Data = structuredData } } messaging.policyManager.ManagePolicy(config.PolicyPayload(payload))