GoMFT/internal/db/middleware/encryption_middleware_test.go

package middleware

import (
	"testing"
	"time"

	"strings"

	"github.com/glebarez/sqlite"
	"github.com/starfleetcptn/gomft/internal/encryption"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

// TestModel is a simple model for testing encryption middleware
type TestModel struct {
	ID                uint      `gorm:"primarykey"`
	Name              string    `gorm:"not null"`
	Password          string    `gorm:"-"` // Not stored in DB, only for form input
	EncryptedPassword string    `gorm:"column:encrypted_password"`
	APIKey            string    `gorm:"-"` // Not stored in DB, only for form input
	EncryptedAPIKey   string    `gorm:"column:encrypted_api_key"`
	CreatedAt         time.Time `gorm:"not null"`
	UpdatedAt         time.Time `gorm:"not null"`
}

// StorageProvider is a simplified version of the real model for testing
type StorageProvider struct {
	ID                    uint      `gorm:"primarykey"`
	Name                  string    `gorm:"not null"`
	Type                  string    `gorm:"not null"`
	Password              string    `gorm:"-"` // Not stored in DB
	EncryptedPassword     string    `gorm:"column:encrypted_password"`
	SecretKey             string    `gorm:"-"` // Not stored in DB
	EncryptedSecretKey    string    `gorm:"column:encrypted_secret_key"`
	ClientSecret          string    `gorm:"-"` // Not stored in DB
	EncryptedClientSecret string    `gorm:"column:encrypted_client_secret"`
	RefreshToken          string    `gorm:"-"` // Not stored in DB
	EncryptedRefreshToken string    `gorm:"column:encrypted_refresh_token"`
	CreatedAt             time.Time `gorm:"not null"`
	UpdatedAt             time.Time `gorm:"not null"`
}

// GetSensitiveFields returns a map of field names to values that need encryption
func (sp *StorageProvider) GetSensitiveFields() map[string]string {
	sensitiveFields := make(map[string]string)

	if sp.Password != "" {
		sensitiveFields["Password"] = sp.Password
	}
	if sp.SecretKey != "" {
		sensitiveFields["SecretKey"] = sp.SecretKey
	}
	if sp.ClientSecret != "" {
		sensitiveFields["ClientSecret"] = sp.ClientSecret
	}
	if sp.RefreshToken != "" {
		sensitiveFields["RefreshToken"] = sp.RefreshToken
	}

	return sensitiveFields
}

func setupTestDB(t *testing.T) *gorm.DB {
	// Initialize in-memory SQLite database
	db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
	require.NoError(t, err, "Failed to connect to in-memory database")

	// Migrate the test models
	err = db.AutoMigrate(&TestModel{}, &StorageProvider{})
	require.NoError(t, err, "Failed to migrate test models")

	return db
}

func setupEncryptionMiddleware(t *testing.T) (*EncryptionMiddleware, error) {
	// Initialize the encryption key manager for testing
	err := encryption.InitializeKeyManager("test-key")
	require.NoError(t, err, "Failed to initialize key manager")

	// Create the encryption middleware
	return NewEncryptionMiddleware()
}

func TestEncryptionMiddlewareWithGenericModel(t *testing.T) {
	// Setup
	db := setupTestDB(t)
	middleware, err := setupEncryptionMiddleware(t)
	require.NoError(t, err, "Failed to setup encryption middleware")

	// Register hooks with GORM
	middleware.RegisterHooks(db)

	// Create a test model
	testModel := &TestModel{
		Name:     "Test User",
		Password: "securePassword123",
		APIKey:   "api-key-12345",
	}

	// Save the model - should trigger encryption
	err = db.Create(testModel).Error
	require.NoError(t, err, "Failed to save test model")

	// Verify encrypted fields are set and original fields are cleared
	assert.Empty(t, testModel.Password, "Password should be cleared after save")
	assert.Empty(t, testModel.APIKey, "APIKey should be cleared after save")
	assert.NotEmpty(t, testModel.EncryptedPassword, "EncryptedPassword should be set")
	assert.NotEmpty(t, testModel.EncryptedAPIKey, "EncryptedAPIKey should be set")
	assert.True(t, strings.HasPrefix(testModel.EncryptedPassword, encryption.EncryptedPrefix), "EncryptedPassword should have encryption prefix")
	assert.True(t, strings.HasPrefix(testModel.EncryptedAPIKey, encryption.EncryptedPrefix), "EncryptedAPIKey should have encryption prefix")

	// Test retrieval and automatic decryption
	retrievedModel := new(TestModel)
	err = db.First(retrievedModel, testModel.ID).Error
	require.NoError(t, err, "Failed to retrieve test model")

	// Verify decryption
	assert.Equal(t, "securePassword123", retrievedModel.Password, "Password should be automatically decrypted")
	assert.Equal(t, "api-key-12345", retrievedModel.APIKey, "APIKey should be automatically decrypted")
	assert.NotEmpty(t, retrievedModel.EncryptedPassword, "EncryptedPassword should remain set")
	assert.NotEmpty(t, retrievedModel.EncryptedAPIKey, "EncryptedAPIKey should remain set")
}

func TestEncryptionMiddlewareWithStorageProvider(t *testing.T) {
	// Setup
	db := setupTestDB(t)
	middleware, err := setupEncryptionMiddleware(t)
	require.NoError(t, err, "Failed to setup encryption middleware")

	// Register hooks with GORM
	middleware.RegisterHooks(db)

	// Create a storage provider
	provider := &StorageProvider{
		Name:         "Test S3",
		Type:         "s3",
		Password:     "testPassword",
		SecretKey:    "testSecretKey",
		ClientSecret: "testClientSecret",
		RefreshToken: "testRefreshToken",
	}

	// Save the provider - should trigger encryption
	err = db.Create(provider).Error
	require.NoError(t, err, "Failed to save storage provider")

	// Verify encrypted fields are set and original fields are cleared
	assert.Empty(t, provider.Password, "Password should be cleared after save")
	assert.Empty(t, provider.SecretKey, "SecretKey should be cleared after save")
	assert.Empty(t, provider.ClientSecret, "ClientSecret should be cleared after save")
	assert.Empty(t, provider.RefreshToken, "RefreshToken should be cleared after save")
	assert.NotEmpty(t, provider.EncryptedPassword, "EncryptedPassword should be set")
	assert.NotEmpty(t, provider.EncryptedSecretKey, "EncryptedSecretKey should be set")
	assert.NotEmpty(t, provider.EncryptedClientSecret, "EncryptedClientSecret should be set")
	assert.NotEmpty(t, provider.EncryptedRefreshToken, "EncryptedRefreshToken should be set")

	// Test retrieval and automatic decryption
	retrievedProvider := new(StorageProvider)
	err = db.First(retrievedProvider, provider.ID).Error
	require.NoError(t, err, "Failed to retrieve storage provider")

	// Verify decryption
	assert.Equal(t, "testPassword", retrievedProvider.Password, "Password should be automatically decrypted")
	assert.Equal(t, "testSecretKey", retrievedProvider.SecretKey, "SecretKey should be automatically decrypted")
	assert.Equal(t, "testClientSecret", retrievedProvider.ClientSecret, "ClientSecret should be automatically decrypted")
	assert.Equal(t, "testRefreshToken", retrievedProvider.RefreshToken, "RefreshToken should be automatically decrypted")
}

func TestEncryptionMiddlewareWithMultipleRecords(t *testing.T) {
	// Setup
	db := setupTestDB(t)
	middleware, err := setupEncryptionMiddleware(t)
	require.NoError(t, err, "Failed to setup encryption middleware")

	// Register hooks with GORM
	middleware.RegisterHooks(db)

	// Create multiple test models
	models := []TestModel{
		{Name: "User 1", Password: "password1", APIKey: "apikey1"},
		{Name: "User 2", Password: "password2", APIKey: "apikey2"},
		{Name: "User 3", Password: "password3", APIKey: "apikey3"},
	}

	// Save all models
	err = db.Create(&models).Error
	require.NoError(t, err, "Failed to save multiple test models")

	// Retrieve all models
	var retrievedModels []TestModel
	err = db.Find(&retrievedModels).Error
	require.NoError(t, err, "Failed to retrieve all test models")

	// Verify count
	assert.Equal(t, 3, len(retrievedModels), "Should retrieve 3 models")

	// Verify each model was properly decrypted
	expectedPasswords := []string{"password1", "password2", "password3"}
	expectedAPIKeys := []string{"apikey1", "apikey2", "apikey3"}

	for i, model := range retrievedModels {
		assert.Equal(t, expectedPasswords[i], model.Password, "Password should be automatically decrypted")
		assert.Equal(t, expectedAPIKeys[i], model.APIKey, "APIKey should be automatically decrypted")
		assert.NotEmpty(t, model.EncryptedPassword, "EncryptedPassword should remain set")
		assert.NotEmpty(t, model.EncryptedAPIKey, "EncryptedAPIKey should remain set")
	}
}

func TestEncryptionMiddlewareDisabled(t *testing.T) {
	// Setup
	db := setupTestDB(t)
	middleware, err := setupEncryptionMiddleware(t)
	require.NoError(t, err, "Failed to setup encryption middleware")

	// Register hooks with GORM
	middleware.RegisterHooks(db)

	// Disable the middleware
	middleware.Disable()
	assert.False(t, middleware.IsEnabled(), "Middleware should be disabled")

	// Create a test model
	testModel := &TestModel{
		Name:     "Test User",
		Password: "securePassword123",
		APIKey:   "api-key-12345",
	}

	// Save the model - should NOT trigger encryption since middleware is disabled
	err = db.Create(testModel).Error
	require.NoError(t, err, "Failed to save test model")

	// Verify sensitive fields are NOT encrypted
	assert.Equal(t, "securePassword123", testModel.Password, "Password should not be cleared when middleware is disabled")
	assert.Equal(t, "api-key-12345", testModel.APIKey, "APIKey should not be cleared when middleware is disabled")
	assert.Empty(t, testModel.EncryptedPassword, "EncryptedPassword should not be set when middleware is disabled")
	assert.Empty(t, testModel.EncryptedAPIKey, "EncryptedAPIKey should not be set when middleware is disabled")

	// Re-enable the middleware for subsequent operations
	middleware.Enable()
	assert.True(t, middleware.IsEnabled(), "Middleware should be enabled")

	// Update the model - should now trigger encryption
	testModel.Password = "newPassword456"
	err = db.Save(testModel).Error
	require.NoError(t, err, "Failed to update test model")

	// Verify encryption now happened
	assert.Empty(t, testModel.Password, "Password should be cleared after update with middleware enabled")
	assert.NotEmpty(t, testModel.EncryptedPassword, "EncryptedPassword should be set after update with middleware enabled")
}