package crypto

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/des"
	"crypto/sha1"
	"fmt"
)

// AESCBCEncrypt encrypts data using AES-CBC mode with PKCS5 padding.
// Supports all AES key sizes: 16 bytes (AES-128), 24 bytes (AES-192), or 32 bytes (AES-256).
func AESCBCEncrypt(key, iv, plaintext []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return cbcEncrypt(block, iv, plaintext)
}

// AESCBCDecrypt decrypts data using AES-CBC mode with PKCS5 unpadding.
// Supports all AES key sizes: 16 bytes (AES-128), 24 bytes (AES-192), or 32 bytes (AES-256).
func AESCBCDecrypt(key, iv, ciphertext []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return cbcDecrypt(block, iv, ciphertext)
}

// DES3Encrypt encrypts data using 3DES-CBC mode with PKCS5 padding.
func DES3Encrypt(key, iv, plaintext []byte) ([]byte, error) {
	block, err := des.NewTripleDESCipher(key)
	if err != nil {
		return nil, err
	}
	return cbcEncrypt(block, iv, plaintext)
}

// DES3Decrypt decrypts data using 3DES-CBC mode with PKCS5 unpadding.
func DES3Decrypt(key, iv, ciphertext []byte) ([]byte, error) {
	block, err := des.NewTripleDESCipher(key)
	if err != nil {
		return nil, err
	}
	return cbcDecrypt(block, iv, ciphertext)
}

// gcmNonceSize is the AES-GCM standard nonce size used by Chromium's v10/v20
// cipher formats. Cross-platform because the v20 ciphertext layout is the
// same regardless of host OS (only Windows currently produces v20).
const gcmNonceSize = 12

// chromiumCBCIV is the fixed IV Chromium uses for AES-CBC v10/v11 (macOS/Linux).
var chromiumCBCIV = bytes.Repeat([]byte{0x20}, aes.BlockSize)

// kEmptyKey is Chromium's decrypt-only fallback for data corrupted by a KWallet
// race in Chrome ~89 (crbug.com/40055416). Matches kEmptyKey in os_crypt_linux.cc.
var kEmptyKey = PBKDF2Key([]byte(""), []byte("saltysalt"), 1, 16, sha1.New)

// DecryptChromiumGCM decrypts a prefixed AES-GCM blob: version(3B)+nonce(12B)+ct+tag.
// Used by Windows v10 (AES-256) and v20; the layout is identical and platform-neutral.
func DecryptChromiumGCM(key, ciphertext []byte) ([]byte, error) {
	if len(ciphertext) < versionPrefixLen+gcmNonceSize {
		return nil, errShortCiphertext
	}
	nonce := ciphertext[versionPrefixLen : versionPrefixLen+gcmNonceSize]
	payload := ciphertext[versionPrefixLen+gcmNonceSize:]
	return AESGCMDecrypt(key, nonce, payload)
}

// DecryptChromiumCBC decrypts a prefixed AES-CBC blob (version(3B)+ct) with Chromium's
// fixed IV, retrying with kEmptyKey to recover crbug.com/40055416 KWallet-corrupted data.
// Used by macOS/Linux v10 and Linux v11 (both AES-128).
func DecryptChromiumCBC(key, ciphertext []byte) ([]byte, error) {
	if len(ciphertext) < versionPrefixLen+aes.BlockSize {
		return nil, errShortCiphertext
	}
	payload := ciphertext[versionPrefixLen:]
	plaintext, err := AESCBCDecrypt(key, chromiumCBCIV, payload)
	if err == nil {
		return plaintext, nil
	}
	if alt, altErr := AESCBCDecrypt(kEmptyKey, chromiumCBCIV, payload); altErr == nil {
		return alt, nil
	}
	return nil, err
}

// AESGCMEncrypt encrypts data using AES-GCM mode.
func AESGCMEncrypt(key, nonce, plaintext []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	if len(nonce) != aead.NonceSize() {
		return nil, errInvalidNonceLen
	}
	return aead.Seal(nil, nonce, plaintext, nil), nil
}

// AESGCMDecrypt decrypts data using AES-GCM mode.
func AESGCMDecrypt(key, nonce, ciphertext []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	if len(nonce) != aead.NonceSize() {
		return nil, errInvalidNonceLen
	}
	return aead.Open(nil, nonce, ciphertext, nil)
}

// AESGCMDecryptBlob decrypts a blob shaped as [12B nonce][ciphertext+16B GCM tag] with caller-supplied AAD.
// Used by protocols that wrap AES-GCM output with a fixed-length nonce prefix (Yandex passwords/cards).
func AESGCMDecryptBlob(key, blob, aad []byte) ([]byte, error) {
	if len(blob) < gcmNonceSize {
		return nil, errShortCiphertext
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	return aead.Open(nil, blob[:gcmNonceSize], blob[gcmNonceSize:], aad)
}

// cbcEncrypt adds PKCS5 padding and encrypts plaintext in CBC mode.
func cbcEncrypt(block cipher.Block, iv, plaintext []byte) ([]byte, error) {
	if len(iv) != block.BlockSize() {
		return nil, errInvalidIVLength
	}

	padded := pkcs5Padding(plaintext, block.BlockSize())
	dst := make([]byte, len(padded))
	cipher.NewCBCEncrypter(block, iv).CryptBlocks(dst, padded)
	return dst, nil
}

// cbcDecrypt decrypts ciphertext in CBC mode and removes PKCS5 padding.
func cbcDecrypt(block cipher.Block, iv, ciphertext []byte) ([]byte, error) {
	bs := block.BlockSize()
	if len(iv) != bs {
		return nil, errInvalidIVLength
	}
	if len(ciphertext) < bs {
		return nil, errShortCiphertext
	}
	if len(ciphertext)%bs != 0 {
		return nil, errInvalidBlockSize
	}

	dst := make([]byte, len(ciphertext))
	cipher.NewCBCDecrypter(block, iv).CryptBlocks(dst, ciphertext)

	dst, err := pkcs5UnPadding(dst, bs)
	if err != nil {
		return nil, fmt.Errorf("decrypt: %w", err)
	}
	return dst, nil
}

// paddingZero pads src with zero bytes to the given length.
// Returns src unchanged if already long enough; otherwise returns a new slice.
func paddingZero(src []byte, length int) []byte {
	if len(src) >= length {
		return src
	}
	dst := make([]byte, length)
	copy(dst, src)
	return dst
}

// pkcs5Padding adds PKCS5/PKCS7 padding to src.
// Always returns a new slice; never modifies src.
func pkcs5Padding(src []byte, blockSize int) []byte {
	n := blockSize - (len(src) % blockSize)
	dst := make([]byte, len(src)+n)
	copy(dst, src)
	for i := len(src); i < len(dst); i++ {
		dst[i] = byte(n)
	}
	return dst
}

// pkcs5UnPadding removes PKCS5/PKCS7 padding from src.
func pkcs5UnPadding(src []byte, blockSize int) ([]byte, error) {
	length := len(src)
	if length == 0 {
		return nil, errInvalidPadding
	}
	padding := int(src[length-1])
	if padding < 1 || padding > blockSize || padding > length {
		return nil, errInvalidPadding
	}
	for _, b := range src[length-padding:] {
		if int(b) != padding {
			return nil, errInvalidPadding
		}
	}
	return src[:length-padding], nil
}