probo/pkg/agent/guardrail/sensitive_data_test.go

// Copyright (c) 2026 Probo Inc <hello@probo.com>.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package guardrail_test

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.probo.inc/probo/pkg/agent/guardrail"
	"go.probo.inc/probo/pkg/llm"
)

func assistantMessage(text string) llm.Message {
	return llm.Message{
		Role:  llm.RoleAssistant,
		Parts: []llm.Part{llm.TextPart{Text: text}},
	}
}

func TestSensitiveDataGuardrail_Check(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name     string
		text     string
		tripwire bool
	}{
		// Safe messages
		{"safe message", "Here is your compliance overview.", false},
		{"safe message with numbers", "You have 42 controls across 3 frameworks.", false},

		// Slack tokens
		{"slack bot token", "The token is xoxb-1234-abcd", true},
		{"slack user token", "Use xoxp-secret-token to authenticate", true},
		{"slack app token", "App token: xoxa-2-abc", true},
		{"slack session token", "Session: xoxs-abc123", true},
		{"slack app-level token", "Token: xapp-1-abc123", true},

		// GitHub tokens
		{"github personal access token", "Use ghp_abc123def456 for auth", true},
		{"github oauth token", "Token: gho_abc123", true},
		{"github user-to-server token", "Token: ghu_abc123", true},
		{"github server-to-server token", "Token: ghs_abc123", true},
		{"github refresh token", "Refresh: ghr_abc123", true},

		// Cloud provider keys
		{"aws access key", "AWS key: AKIAIOSFODNN7EXAMPLE", true},

		// Payment provider keys
		{"stripe live key", "Stripe key: sk_live_abc123", true},
		{"stripe test key", "Stripe key: sk_test_abc123", true},

		// LLM provider keys
		{"openai key", "The API key is sk-proj-abc123", true},
		{"anthropic key", "Key: sk-ant-api03-abc123", true},
		{"sk prefix not a false positive", "This is a risk-based approach to task-management.", false},

		// JWT tokens
		{"jwt token", "Token: eyJhbGciOiJIUzI1NiJ9.payload.sig", true},

		// Authorization headers
		{"bearer auth", "Authorization: Bearer abc123", true},
		{"basic auth", "Authorization: Basic dXNlcjpwYXNz", true},
		{"bearer auth uppercase", "BEARER token123", true},

		// PEM / certificates
		{"pem private key", "-----BEGIN RSA PRIVATE KEY-----\nMIIE...", true},
		{"pem certificate", "-----BEGIN CERTIFICATE-----\nMIIE...", true},

		// Connection strings
		{"postgres uri", "Connect to postgres://user:pass@host/db", true},
		{"postgresql uri", "Connect to postgresql://user:pass@host/db", true},
		{"mongodb uri", "Use mongodb://user:pass@host/db", true},
		{"mysql uri", "Use mysql://user:pass@host/db", true},
		{"redis uri", "Cache at redis://localhost:6379", true},
		{"amqp uri", "Queue at amqp://guest:guest@host/vhost", true},

		// Generic secret field names
		{"encryption_key", "The encryption_key is set in config", true},
		{"signing_secret", "Your signing_secret was rotated", true},
		{"secret_key", "The secret_key value is abc", true},
		{"private_key", "Set private_key in the env", true},
		{"client_secret", "The client_secret is abc123", true},
		{"access_token", "Use this access_token to call the API", true},
		{"api_key", "Your api_key is xyz", true},
		{"apikey", "Set apikey in headers", true},
		{"password", "Your password is hunter2", true},

		// Raw SQL
		{"sql select", "SELECT * FROM users WHERE id = 1", true},
		{"sql insert", "INSERT INTO users VALUES (1, 'admin')", true},
		{"sql update", "UPDATE users SET name = 'foo'", true},
		{"sql delete", "DELETE FROM users WHERE id = 1", true},
		{"sql drop", "DROP TABLE users", true},
	}

	g := guardrail.NewSensitiveDataGuardrail()

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			result, err := g.Check(context.Background(), assistantMessage(tt.text))

			require.NoError(t, err)
			assert.Equal(t, tt.tripwire, result.Tripwire)
		})
	}
}