teleport/pkg/encryption/encryption.go

package encryption

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"fmt"
	"io"
	"math"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/pbkdf2"
)

const (
	// PBKDF2 parameters for key derivation
	PBKDF2Iterations = 100000 // OWASP recommended minimum
	PBKDF2KeyLength  = 32     // 256 bits
	PBKDF2SaltLength = 16     // 128 bits

	// Replay protection parameters
	MaxPacketAge = 5 * time.Minute // Maximum age for UDP packets
	NonceWindow  = 1000            // Number of nonces to track for replay protection

	// RawKeyPrefix marks a hex-encoded 32-byte AES key that must not go through PBKDF2.
	RawKeyPrefix = "raw:"
	RawKeyHexLen = 64 // 32 bytes
)

// deriveCache memoizes PBKDF2 so handshake/UDP paths never pay 100k iterations twice.
var deriveCache sync.Map // map[string][]byte

// GenerateRawKey returns a raw: prefixed hex encoding of 32 random bytes.
// Paste the entire string into encryption_key; it is used as an AES-256 key (no PBKDF2).
func GenerateRawKey() (string, error) {
	key := make([]byte, PBKDF2KeyLength)
	if _, err := io.ReadFull(rand.Reader, key); err != nil {
		return "", fmt.Errorf("failed to generate random key: %v", err)
	}
	return RawKeyPrefix + hex.EncodeToString(key), nil
}

// IsRawKey reports whether material is a raw 256-bit key (raw: + 64 hex chars).
func IsRawKey(material string) bool {
	return strings.HasPrefix(material, RawKeyPrefix)
}

// ResolveKey returns a 32-byte AES key from config material.
//
//   - raw:<64 hex>: hex-decode, no PBKDF2 (new --generate-key / generated configs)
//   - anything else, including legacy unprefixed 64-hex from older --generate-key:
//     PBKDF2 with the historical password-derived salt (existing configs keep working)
func ResolveKey(material string) ([]byte, error) {
	if strings.HasPrefix(material, RawKeyPrefix) {
		hexStr := material[len(RawKeyPrefix):]
		key, err := hex.DecodeString(hexStr)
		if err != nil {
			return nil, fmt.Errorf("invalid raw encryption key: %w", err)
		}
		if len(key) != PBKDF2KeyLength {
			return nil, fmt.Errorf("raw encryption key must be %d bytes, got %d", PBKDF2KeyLength, len(key))
		}
		return key, nil
	}
	return DeriveKey(material), nil
}

// DeriveKey derives an encryption key from a password using PBKDF2.
// The first call for a given password runs 100k iterations; later calls return the cached key.
func DeriveKey(password string) []byte {
	if v, ok := deriveCache.Load(password); ok {
		return v.([]byte)
	}

	// Historical salt is SHA256(password)[:16]. That is not a random salt; keep it only
	// so existing passphrase and unprefixed hex configs still derive the same key.
	hasher := sha256.New()
	hasher.Write([]byte(password))
	passwordHash := hasher.Sum(nil)

	salt := make([]byte, PBKDF2SaltLength)
	copy(salt, passwordHash[:PBKDF2SaltLength])

	key := pbkdf2.Key([]byte(password), salt, PBKDF2Iterations, PBKDF2KeyLength, sha256.New)
	actual, _ := deriveCache.LoadOrStore(password, key)
	return actual.([]byte)
}

// DeriveKeyWithSalt derives an encryption key from a password using PBKDF2 with a custom salt
func DeriveKeyWithSalt(password string, salt []byte) ([]byte, error) {
	if len(salt) != PBKDF2SaltLength {
		return nil, fmt.Errorf("salt length must be %d bytes, got %d", PBKDF2SaltLength, len(salt))
	}

	key := pbkdf2.Key([]byte(password), salt, PBKDF2Iterations, PBKDF2KeyLength, sha256.New)
	return key, nil
}

// GenerateSalt generates a cryptographically secure random salt
func GenerateSalt() ([]byte, error) {
	salt := make([]byte, PBKDF2SaltLength)
	_, err := io.ReadFull(rand.Reader, salt)
	return salt, err
}

// EncryptData encrypts data using AES-GCM
func EncryptData(data []byte, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, err
	}

	ciphertext := gcm.Seal(nonce, nonce, data, nil)
	return ciphertext, nil
}

// DecryptData decrypts data using AES-GCM
func DecryptData(data []byte, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	nonceSize := gcm.NonceSize()
	if len(data) < nonceSize {
		return nil, fmt.Errorf("ciphertext too short")
	}

	nonce, ciphertext := data[:nonceSize], data[nonceSize:]
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, err
	}

	return plaintext, nil
}

// ReplayProtection tracks nonces to prevent replay attacks
type ReplayProtection struct {
	nonces    map[uint64]time.Time
	mutex     sync.RWMutex
	maxNonces int // Maximum number of nonces to track
}

// NewReplayProtection creates a new replay protection instance
func NewReplayProtection() *ReplayProtection {
	return &ReplayProtection{
		nonces:    make(map[uint64]time.Time),
		maxNonces: NonceWindow * 2, // Allow 2x the window size for safety
	}
}

// IsValidNonce checks if a nonce is valid (not replayed and not too old)
func (rp *ReplayProtection) IsValidNonce(nonce uint64, timestamp int64) bool {
	rp.mutex.Lock()
	defer rp.mutex.Unlock()

	now := time.Now()
	packetTime := time.Unix(timestamp, 0)

	// Check if packet is too old
	if now.Sub(packetTime) > MaxPacketAge {
		return false
	}

	// Check if nonce was already used
	if _, exists := rp.nonces[nonce]; exists {
		return false
	}

	// Add nonce to tracking
	rp.nonces[nonce] = now

	// Clean up old nonces to prevent memory leaks
	if len(rp.nonces) > rp.maxNonces {
		rp.cleanupOldNonces(now)
	}

	return true
}

// cleanupOldNonces removes nonces older than MaxPacketAge
func (rp *ReplayProtection) cleanupOldNonces(now time.Time) {
	for nonce, timestamp := range rp.nonces {
		if now.Sub(timestamp) > MaxPacketAge {
			delete(rp.nonces, nonce)
		}
	}
}

// ValidatePacketTimestamp validates that a packet timestamp is within acceptable range
func ValidatePacketTimestamp(timestamp int64) bool {
	now := time.Now().Unix()
	packetTime := timestamp

	// Check if packet is too old or from the future
	age := now - packetTime
	return age >= 0 && age <= int64(MaxPacketAge.Seconds())
}

// ConstantTimeCompare performs a constant-time comparison of two byte slices
func ConstantTimeCompare(a, b []byte) bool {
	return subtle.ConstantTimeCompare(a, b) == 1
}

// ValidateEncryptionKey validates that an encryption key meets security requirements
func ValidateEncryptionKey(key string) error {
	if strings.HasPrefix(key, RawKeyPrefix) {
		hexStr := key[len(RawKeyPrefix):]
		if len(hexStr) != RawKeyHexLen {
			return fmt.Errorf("raw encryption key must be %d hex characters (32 bytes), got %d", RawKeyHexLen, len(hexStr))
		}
		decoded, err := hex.DecodeString(hexStr)
		if err != nil {
			return fmt.Errorf("raw encryption key is not valid hex: %w", err)
		}
		if len(decoded) != PBKDF2KeyLength {
			return fmt.Errorf("raw encryption key must decode to %d bytes", PBKDF2KeyLength)
		}
		return nil
	}

	if len(key) < 32 {
		return fmt.Errorf("encryption key must be at least 32 characters long")
	}

	// Check for common weak keys
	weakKeys := []string{
		"password", "123456", "admin", "test", "default",
		"your-secure-encryption-key-change-this-to-something-random",
		"test-encryption-key-12345", "teleport-key", "secret",
	}

	for _, weak := range weakKeys {
		if key == weak {
			return fmt.Errorf("encryption key is too weak, please use a strong random key")
		}
	}

	// Check for sufficient entropy (basic check)
	entropy := calculateEntropy(key)
	if entropy < 3.5 { // Minimum entropy threshold
		return fmt.Errorf("encryption key has insufficient entropy, please use a more random key")
	}

	return nil
}

// calculateEntropy calculates the Shannon entropy of a string
func calculateEntropy(s string) float64 {
	freq := make(map[rune]int)
	for _, r := range s {
		freq[r]++
	}

	entropy := 0.0
	length := float64(len(s))

	for _, count := range freq {
		p := float64(count) / length
		entropy -= p * log2(p)
	}

	return entropy
}

// log2 calculates log base 2
func log2(x float64) float64 {
	return math.Log2(x)
}