package kdbx

import (
	"bytes"
	"compress/gzip"
	"crypto/rand"
	"crypto/sha256"
	"crypto/sha512"
	"crypto/subtle"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
)

// Binary is an attachment, stored once in the file's pool and referenced
// from entries by index.
type Binary struct {
	// Flags is the KDBX protection flag byte; bit 0 asks the client to keep
	// the content in protected memory.
	Flags byte
	Data  []byte
}

// Database is a decrypted KDBX file: its format parameters, its attachment
// pool and its XML document.
type Database struct {
	Version     uint32
	Cipher      []byte
	Compression uint32
	KDF         variantDict
	PublicData  variantDict
	Binaries    []Binary
	// Root is the KeePassFile element. Everything the format carries lives
	// under it, including elements this package does not interpret.
	Root *Node
}

// KDFParams describes the work factor a new database is created with.
type KDFParams struct {
	// Argon2 parameters. MemoryKiB and Iterations are what make an attack
	// expensive; they are also what a merging server pays on every push.
	MemoryKiB  uint32
	Iterations uint32
	Lanes      uint8
}

// DefaultKDF is deliberately lighter than the KeePass desktop default: a
// server re-derives the key three times per merge, so the cost is paid on
// every save by every member of the vault.
var DefaultKDF = KDFParams{MemoryKiB: 32 * 1024, Iterations: 2, Lanes: 2}

// Open decrypts and parses a KDBX 4 file.
func Open(data []byte, password string) (*Database, error) {
	if len(data) < 12 {
		return nil, errors.New("kdbx: file is too short")
	}
	if binary.LittleEndian.Uint32(data[0:]) != signature1 ||
		binary.LittleEndian.Uint32(data[4:]) != signature2 {
		return nil, errors.New("kdbx: not a KeePass 2 database")
	}
	version := binary.LittleEndian.Uint32(data[8:])
	if major := version >> 16; major != 4 {
		return nil, unsupported("KDBX %d.%d; this server stores KDBX 4 vaults",
			major, version&0xFFFF)
	}

	db := &Database{Version: version}
	p := 12
	for {
		if p+5 > len(data) {
			return nil, errors.New("kdbx: truncated header")
		}
		id := data[p]
		size := int(binary.LittleEndian.Uint32(data[p+1:]))
		p += 5
		if size < 0 || p+size > len(data) {
			return nil, errors.New("kdbx: truncated header field")
		}
		field := data[p : p+size]
		p += size

		switch id {
		case hdrEnd:
			goto headerDone
		case hdrCipherID:
			db.Cipher = append([]byte(nil), field...)
		case hdrCompression:
			if len(field) != 4 {
				return nil, errors.New("kdbx: malformed compression flags")
			}
			db.Compression = binary.LittleEndian.Uint32(field)
		case hdrMasterSeed:
			// Kept locally; it never leaves this function.
		case hdrEncryptionIV:
		case hdrKdfParameters:
			vd, err := parseVariantDict(field)
			if err != nil {
				return nil, err
			}
			db.KDF = vd
		case hdrPublicCustomData:
			vd, err := parseVariantDict(field)
			if err != nil {
				return nil, err
			}
			db.PublicData = vd
		case hdrComment:
		case hdrTransformSeed, hdrTransformRounds, hdrInnerStreamKey3,
			hdrStreamStartBytes, hdrInnerStreamID3:
			return nil, unsupported("KDBX 3 header field %d in a KDBX 4 file", id)
		}
	}
headerDone:
	headerBytes := data[:p]
	masterSeed := headerField(data[12:p], hdrMasterSeed)
	iv := headerField(data[12:p], hdrEncryptionIV)
	if len(masterSeed) != 32 {
		return nil, errors.New("kdbx: master seed has the wrong length")
	}
	if len(db.Cipher) != 16 {
		return nil, errors.New("kdbx: file names no payload cipher")
	}

	if p+64 > len(data) {
		return nil, errors.New("kdbx: file ends before its integrity hashes")
	}
	storedSHA := data[p : p+32]
	storedHMAC := data[p+32 : p+64]
	sum := sha256.Sum256(headerBytes)
	if subtle.ConstantTimeCompare(sum[:], storedSHA) != 1 {
		return nil, errors.New("kdbx: header does not match its checksum")
	}

	k, err := deriveKeys(masterSeed, db.KDF, password)
	if err != nil {
		return nil, err
	}
	if subtle.ConstantTimeCompare(headerHMAC(k.hmac, headerBytes), storedHMAC) != 1 {
		return nil, ErrWrongPassword
	}

	payload, err := readBlocks(data[p+64:], k.hmac)
	if err != nil {
		return nil, err
	}
	plain, err := decryptPayload(db.Cipher, k.payload, iv, payload)
	if err != nil {
		return nil, err
	}
	if db.Compression == 1 {
		if plain, err = gunzip(plain); err != nil {
			return nil, fmt.Errorf("kdbx: decompress: %w", err)
		}
	} else if db.Compression != 0 {
		return nil, unsupported("compression algorithm %d", db.Compression)
	}

	streamID, streamKey, xmlBytes, err := db.readInnerHeader(plain)
	if err != nil {
		return nil, err
	}
	root, err := parseXML(xmlBytes)
	if err != nil {
		return nil, err
	}
	db.Root = root
	if err := db.applyInnerStream(streamID, streamKey, false); err != nil {
		return nil, err
	}
	return db, nil
}

// Marshal encrypts the database back into a KDBX 4 file. Every save gets
// fresh seeds, a fresh IV and a fresh inner stream key.
func (d *Database) Marshal(password string) ([]byte, error) {
	if d.Root == nil {
		return nil, errors.New("kdbx: database has no document")
	}
	cipherID := d.Cipher
	if len(cipherID) != 16 {
		cipherID = cipherAES256CBC
	}
	ivLen := 16
	if bytes.Equal(cipherID, cipherChaCha20Outer) {
		ivLen = 12
	}
	masterSeed := randomBytes(32)
	iv := randomBytes(ivLen)
	streamKey := randomBytes(64)

	kdf := d.KDF
	if _, ok := kdf.get("$UUID"); !ok {
		kdf = newArgon2Params(DefaultKDF)
	}
	// A new salt per save, so two saves never share a transformed key.
	switch {
	case bytes.Equal(kdf.bytes("$UUID"), kdfAES):
		kdf.set(vdByteArray, "S", randomBytes(32))
	default:
		kdf.set(vdByteArray, "S", randomBytes(32))
	}
	d.KDF = kdf

	// Protect values on a copy so the in-memory tree stays plaintext and can
	// be marshalled again.
	protectedRoot := d.Root.clone()
	if err := applyStream(protectedRoot, streamChaCha20, streamKey, true); err != nil {
		return nil, err
	}

	var payload bytes.Buffer
	payload.Write(d.innerHeader(streamChaCha20, streamKey))
	payload.Write(marshalXML(protectedRoot))

	body := payload.Bytes()
	if d.Compression == 1 {
		var gz bytes.Buffer
		w := gzip.NewWriter(&gz)
		if _, err := w.Write(body); err != nil {
			return nil, err
		}
		if err := w.Close(); err != nil {
			return nil, err
		}
		body = gz.Bytes()
	}

	version := d.Version
	if version>>16 != 4 {
		version = Version4
	}
	header := d.outerHeader(version, cipherID, masterSeed, iv, kdf)

	k, err := deriveKeys(masterSeed, kdf, password)
	if err != nil {
		return nil, err
	}
	encrypted, err := encryptPayload(cipherID, k.payload, iv, body)
	if err != nil {
		return nil, err
	}

	out := make([]byte, 0, len(header)+64+len(encrypted)+64)
	out = append(out, header...)
	sum := sha256.Sum256(header)
	out = append(out, sum[:]...)
	out = append(out, headerHMAC(k.hmac, header)...)
	return append(out, writeBlocks(encrypted, k.hmac)...), nil
}

// outerHeader serializes the unencrypted header.
func (d *Database) outerHeader(version uint32, cipherID, masterSeed, iv []byte,
	kdf variantDict) []byte {
	out := make([]byte, 12)
	binary.LittleEndian.PutUint32(out[0:], signature1)
	binary.LittleEndian.PutUint32(out[4:], signature2)
	binary.LittleEndian.PutUint32(out[8:], version)

	add := func(id byte, data []byte) {
		out = append(out, id)
		out = binary.LittleEndian.AppendUint32(out, uint32(len(data)))
		out = append(out, data...)
	}
	add(hdrCipherID, cipherID)
	add(hdrCompression, binary.LittleEndian.AppendUint32(nil, d.Compression))
	add(hdrMasterSeed, masterSeed)
	add(hdrEncryptionIV, iv)
	add(hdrKdfParameters, kdf.marshal())
	if len(d.PublicData.items) > 0 {
		add(hdrPublicCustomData, d.PublicData.marshal())
	}
	add(hdrEnd, []byte{0x0D, 0x0A, 0x0D, 0x0A})
	return out
}

// readInnerHeader splits the decrypted payload into its inner header, which
// carries the attachment pool, and the XML that follows it.
func (d *Database) readInnerHeader(plain []byte) (streamID uint32, streamKey, xmlBytes []byte, err error) {
	p := 0
	d.Binaries = nil
	for {
		if p+5 > len(plain) {
			return 0, nil, nil, errors.New("kdbx: truncated inner header")
		}
		id := plain[p]
		size := int(binary.LittleEndian.Uint32(plain[p+1:]))
		p += 5
		if size < 0 || p+size > len(plain) {
			return 0, nil, nil, errors.New("kdbx: truncated inner header field")
		}
		field := plain[p : p+size]
		p += size
		switch id {
		case innerEnd:
			return streamID, streamKey, plain[p:], nil
		case innerStreamID:
			if len(field) != 4 {
				return 0, nil, nil, errors.New("kdbx: malformed inner stream id")
			}
			streamID = binary.LittleEndian.Uint32(field)
		case innerStreamKey:
			streamKey = append([]byte(nil), field...)
		case innerBinary:
			if len(field) < 1 {
				return 0, nil, nil, errors.New("kdbx: malformed attachment")
			}
			d.Binaries = append(d.Binaries, Binary{
				Flags: field[0],
				Data:  append([]byte(nil), field[1:]...),
			})
		}
	}
}

func (d *Database) innerHeader(streamID uint32, streamKey []byte) []byte {
	out := []byte{}
	add := func(id byte, data []byte) {
		out = append(out, id)
		out = binary.LittleEndian.AppendUint32(out, uint32(len(data)))
		out = append(out, data...)
	}
	add(innerStreamID, binary.LittleEndian.AppendUint32(nil, streamID))
	add(innerStreamKey, streamKey)
	for _, b := range d.Binaries {
		add(innerBinary, append([]byte{b.Flags}, b.Data...))
	}
	add(innerEnd, nil)
	return out
}

// applyInnerStream decrypts the protected values of this database's tree.
func (d *Database) applyInnerStream(streamID uint32, key []byte, protect bool) error {
	return applyStream(d.Root, streamID, key, protect)
}

// applyStream runs the inner stream cipher across every protected value in
// document order, which is the order KeePass itself uses.
func applyStream(root *Node, streamID uint32, key []byte, protect bool) error {
	var stream *chacha20
	switch streamID {
	case streamNone:
	case streamChaCha20:
		h := sha512.Sum512(key)
		stream = newChaCha20(h[0:32], h[32:44], 0)
	case streamSalsa20:
		return unsupported("Salsa20 inner stream (KDBX 3)")
	default:
		return unsupported("inner stream cipher %d", streamID)
	}

	var err error
	root.walk(func(n *Node) {
		if err != nil || !n.Protected || len(n.Nodes) > 0 {
			return
		}
		if stream == nil {
			return // nothing to do, values are stored in the clear
		}
		if protect {
			plain := []byte(n.Text)
			out := make([]byte, len(plain))
			stream.XORKeyStream(out, plain)
			n.Text = base64.StdEncoding.EncodeToString(out)
			return
		}
		raw, derr := base64.StdEncoding.DecodeString(n.Text)
		if derr != nil {
			err = fmt.Errorf("kdbx: protected value in <%s> is not base64: %w", n.Name, derr)
			return
		}
		out := make([]byte, len(raw))
		stream.XORKeyStream(out, raw)
		n.Text = string(out)
	})
	return err
}

func newArgon2Params(p KDFParams) variantDict {
	if p.MemoryKiB == 0 {
		p = DefaultKDF
	}
	var vd variantDict
	vd.set(vdByteArray, "$UUID", kdfArgon2d)
	vd.set(vdByteArray, "S", randomBytes(32))
	vd.setUint32("P", uint32(max(p.Lanes, 1)))
	vd.setUint64("M", uint64(p.MemoryKiB)*1024)
	vd.setUint64("I", uint64(max(p.Iterations, 1)))
	vd.setUint32("V", argon2Version)
	return vd
}

// headerField re-reads one field out of the raw header, used for values that
// are needed as bytes but not kept on the Database.
func headerField(fields []byte, want byte) []byte {
	p := 0
	for p+5 <= len(fields) {
		id := fields[p]
		size := int(binary.LittleEndian.Uint32(fields[p+1:]))
		p += 5
		if size < 0 || p+size > len(fields) {
			return nil
		}
		if id == want {
			return fields[p : p+size]
		}
		p += size
		if id == hdrEnd {
			return nil
		}
	}
	return nil
}

func gunzip(b []byte) ([]byte, error) {
	r, err := gzip.NewReader(bytes.NewReader(b))
	if err != nil {
		return nil, err
	}
	defer r.Close()
	// A KDBX payload is a password database, not an archive: cap it.
	return io.ReadAll(io.LimitReader(r, 512<<20))
}

func randomBytes(n int) []byte {
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		panic("kdbx: " + err.Error())
	}
	return b
}