uzdb/internal/services/encryption.go

package services

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io"
	"os"
	"sync"

	"go.uber.org/zap"

	"uzdb/internal/config"
	"uzdb/internal/models"
	"uzdb/internal/utils"
)

// EncryptionService handles encryption and decryption of sensitive data
type EncryptionService struct {
	key     []byte
	cipher  cipher.Block
	gcm     cipher.AEAD
	mu      sync.RWMutex
}

var (
	encryptionInstance *EncryptionService
	encryptionOnce     sync.Once
)

// GetEncryptionService returns the singleton encryption service instance
func GetEncryptionService() *EncryptionService {
	return encryptionInstance
}

// InitEncryptionService initializes the encryption service
func InitEncryptionService(cfg *config.EncryptionConfig) (*EncryptionService, error) {
	var err error
	encryptionOnce.Do(func() {
		encryptionInstance = &EncryptionService{}
		err = encryptionInstance.init(cfg)
	})
	return encryptionInstance, err
}

// init initializes the encryption service with a key
func (s *EncryptionService) init(cfg *config.EncryptionConfig) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Try to load existing key or generate new one
	key, err := s.loadOrGenerateKey(cfg)
	if err != nil {
		return fmt.Errorf("failed to load/generate key: %w", err)
	}

	s.key = key

	// Create AES cipher
	block, err := aes.NewCipher(key)
	if err != nil {
		return fmt.Errorf("failed to create cipher: %w", err)
	}
	s.cipher = block

	// Create GCM mode
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return fmt.Errorf("failed to create GCM: %w", err)
	}
	s.gcm = gcm

	config.GetLogger().Info("encryption service initialized")
	return nil
}

// loadOrGenerateKey loads existing key or generates a new one
func (s *EncryptionService) loadOrGenerateKey(cfg *config.EncryptionConfig) ([]byte, error) {
	// Use provided key if available
	if cfg.Key != "" {
		key := []byte(cfg.Key)
		// Ensure key is correct length (32 bytes for AES-256)
		if len(key) < 32 {
			// Pad key
			padded := make([]byte, 32)
			copy(padded, key)
			key = padded
		} else if len(key) > 32 {
			key = key[:32]
		}
		return key, nil
	}

	// Try to load from file
	if cfg.KeyFile != "" {
		if data, err := os.ReadFile(cfg.KeyFile); err == nil {
			key := []byte(data)
			if len(key) >= 32 {
				return key[:32], nil
			}
		}
	}

	// Generate new key
	key := make([]byte, 32) // AES-256
	if _, err := io.ReadFull(rand.Reader, key); err != nil {
		return nil, fmt.Errorf("failed to generate key: %w", err)
	}

	// Save key to file if path provided
	if cfg.KeyFile != "" {
		if err := os.WriteFile(cfg.KeyFile, key, 0600); err != nil {
			config.GetLogger().Warn("failed to save encryption key", zap.Error(err))
		} else {
			config.GetLogger().Info("encryption key generated and saved",
				zap.String("path", cfg.KeyFile))
		}
	}

	return key, nil
}

// Encrypt encrypts plaintext using AES-GCM
func (s *EncryptionService) Encrypt(plaintext string) (string, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.gcm == nil {
		return "", models.ErrEncryptionFailed
	}

	// Generate nonce
	nonce := make([]byte, s.gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", fmt.Errorf("failed to generate nonce: %w", err)
	}

	// Encrypt
	ciphertext := s.gcm.Seal(nonce, nonce, []byte(plaintext), nil)

	// Encode to base64
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// Decrypt decrypts ciphertext using AES-GCM
func (s *EncryptionService) Decrypt(ciphertext string) (string, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.gcm == nil {
		return "", models.ErrEncryptionFailed
	}

	// Decode from base64
	data, err := base64.StdEncoding.DecodeString(ciphertext)
	if err != nil {
		return "", fmt.Errorf("failed to decode ciphertext: %w", err)
	}

	// Verify nonce size
	nonceSize := s.gcm.NonceSize()
	if len(data) < nonceSize {
		return "", models.ErrEncryptionFailed
	}

	// Extract nonce and ciphertext
	nonce, ciphertextBytes := data[:nonceSize], data[nonceSize:]

	// Decrypt
	plaintext, err := s.gcm.Open(nil, nonce, ciphertextBytes, nil)
	if err != nil {
		return "", fmt.Errorf("failed to decrypt: %w", err)
	}

	return string(plaintext), nil
}

// EncryptPassword encrypts a password for storage
func (s *EncryptionService) EncryptPassword(password string) (string, error) {
	if password == "" {
		return "", nil
	}
	return s.Encrypt(password)
}

// DecryptPassword decrypts a stored password
func (s *EncryptionService) DecryptPassword(encryptedPassword string) (string, error) {
	if encryptedPassword == "" {
		return "", nil
	}
	return s.Decrypt(encryptedPassword)
}

// MaskPasswordForLogging masks password for safe logging
func (s *EncryptionService) MaskPasswordForLogging(password string) string {
	return utils.MaskPassword(password)
}