package main

// Bitwarden-compatible crypto for the KNELSecretsManager CLI (Vaultwarden API).
//   - master key: PBKDF2-SHA256(password, email, iterations, 32B) or Argon2id
//   - auth: masterPasswordHash = base64(PBKDF2-SHA256(masterKey, password, 1, 32B))
//   - stretched master key: HKDF-SHA256 expand, info "enc"/"mac" (32B each)
//   - encString "2.iv|ct|mac": AES-256-CBC(encKey32) + HMAC-SHA256(macKey32, iv||ct)
//     ("AesCbc128_HmacSha256_B64" is Bitwarden's legacy misnomer; keys are 32+32 bytes)

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"errors"
	"fmt"
	"strconv"
	"strings"

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

type EncString struct {
	Type byte
	IV   []byte
	CT   []byte
	MAC  []byte
}

func ParseEncString(s string) (*EncString, error) {
	if s == "" {
		return nil, errors.New("empty encstring")
	}
	parts := strings.SplitN(s, ".", 2)
	if len(parts) != 2 {
		return nil, errors.New("encstring missing type header")
	}
	t, err := strconv.Atoi(parts[0])
	if err != nil {
		return nil, fmt.Errorf("bad encstring type: %w", err)
	}
	es := &EncString{Type: byte(t)}
	switch t {
	case 0: // AesCbc256_B64 (legacy, no mac)
		b, err := base64.StdEncoding.DecodeString(parts[1])
		if err != nil {
			return nil, err
		}
		es.IV, es.CT = b[:16], b[16:]
	case 1, 2: // both handled identically on the wire (enc32+mac32 keys)
		seg := strings.Split(parts[1], "|")
		if len(seg) != 3 {
			return nil, errors.New("encstring needs iv|ct|mac")
		}
		if es.IV, err = base64.StdEncoding.DecodeString(seg[0]); err != nil {
			return nil, err
		}
		if es.CT, err = base64.StdEncoding.DecodeString(seg[1]); err != nil {
			return nil, err
		}
		if es.MAC, err = base64.StdEncoding.DecodeString(seg[2]); err != nil {
			return nil, err
		}
	default:
		return nil, fmt.Errorf("unsupported encstring type %d", t)
	}
	return es, nil
}

func (e *EncString) String() string {
	if e.Type == 0 {
		return "0." + base64.StdEncoding.EncodeToString(append(append([]byte{}, e.IV...), e.CT...))
	}
	return fmt.Sprintf("%d.%s|%s|%s", e.Type,
		base64.StdEncoding.EncodeToString(e.IV),
		base64.StdEncoding.EncodeToString(e.CT),
		base64.StdEncoding.EncodeToString(e.MAC))
}

func deriveMasterKey(password, email string, kdfType int, iterations, memory, parallelism uint32) []byte {
	salt := []byte(strings.ToLower(strings.TrimSpace(email)))
	switch kdfType {
	case 1:
		return argon2.IDKey([]byte(password), salt, iterations, memory, uint8(parallelism), 32)
	default:
		return pbkdf2.Key([]byte(password), salt, int(iterations), 32, sha256.New)
	}
}

func masterPasswordHash(masterKey []byte, password string) string {
	return base64.StdEncoding.EncodeToString(pbkdf2.Key(masterKey, []byte(password), 1, 32, sha256.New))
}

// stretchKey expands the 32B master key to a 64B symmetric key (enc 32 | mac 32).
func stretchKey(masterKey []byte) []byte {
	out := make([]byte, 64)
	copy(out[:32], hkdfExpand(masterKey, []byte("enc"), 32))
	copy(out[32:], hkdfExpand(masterKey, []byte("mac"), 32))
	return out
}

func hkdfExpand(key, info []byte, length int) []byte {
	out := make([]byte, 0, length)
	t := []byte{}
	var i byte
	for len(out) < length {
		i++
		h := hmac.New(sha256.New, key)
		h.Write(t)
		h.Write(info)
		h.Write([]byte{i})
		t = h.Sum(nil)
		out = append(out, t...)
	}
	return out[:length]
}

// symDecrypt decrypts a type-1/2 encstring with a 64-byte key (enc32|mac32).
func symDecrypt(key64 []byte, es *EncString) ([]byte, error) {
	if len(key64) != 64 {
		return nil, errors.New("symmetric key must be 64 bytes")
	}
	if es.Type == 0 {
		return aesCBCDecrypt(key64[:32], es.IV, es.CT)
	}
	if es.Type != 1 && es.Type != 2 {
		return nil, fmt.Errorf("unsupported encstring type %d", es.Type)
	}
	mac := hmac.New(sha256.New, key64[32:])
	mac.Write(es.IV)
	mac.Write(es.CT)
	if subtle.ConstantTimeCompare(mac.Sum(nil), es.MAC) != 1 {
		return nil, errors.New("mac mismatch")
	}
	return aesCBCDecrypt(key64[:32], es.IV, es.CT)
}

// symEncrypt encrypts plaintext into a type-2 encstring with a 64-byte key.
func symEncrypt(key64 []byte, plaintext []byte) (string, error) {
	if len(key64) != 64 {
		return "", errors.New("symmetric key must be 64 bytes")
	}
	iv := make([]byte, 16)
	if _, err := rand.Read(iv); err != nil {
		return "", err
	}
	ct, err := aesCBCEncrypt(key64[:32], iv, plaintext)
	if err != nil {
		return "", err
	}
	mac := hmac.New(sha256.New, key64[32:])
	mac.Write(iv)
	mac.Write(ct)
	es := &EncString{Type: 2, IV: iv, CT: ct, MAC: mac.Sum(nil)}
	return es.String(), nil
}

func aesCBCDecrypt(key, iv, ct []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	if len(ct) == 0 || len(ct)%aes.BlockSize != 0 {
		return nil, errors.New("ciphertext not block aligned")
	}
	pt := make([]byte, len(ct))
	cipher.NewCBCDecrypter(block, iv).CryptBlocks(pt, ct)
	return unpad(pt)
}

func aesCBCEncrypt(key, iv, pt []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	ct := make([]byte, len(pad(pt)))
	cipher.NewCBCEncrypter(block, iv).CryptBlocks(ct, pad(pt))
	return ct, nil
}

func unpad(b []byte) ([]byte, error) {
	if len(b) == 0 {
		return nil, errors.New("empty plaintext")
	}
	n := int(b[len(b)-1])
	if n == 0 || n > aes.BlockSize || n > len(b) {
		return nil, errors.New("bad padding")
	}
	return b[:len(b)-n], nil
}

func pad(b []byte) []byte {
	n := aes.BlockSize - len(b)%aes.BlockSize
	out := make([]byte, len(b)+n)
	copy(out, b)
	for i := len(b); i < len(out); i++ {
		out[i] = byte(n)
	}
	return out
}