// Copyright (C) 2022 Marius Schellenberger
// crypto parts taken from https://github.com/filosottile/age

package crypto

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha256"
	"io"

	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/hkdf"
)

const (
	KDFNonceSize = 16
	NonceSize    = chacha20poly1305.NonceSize
	Overhead     = chacha20poly1305.Overhead
	ChunkSize    = 64 * 1024
	FullSize     = ChunkSize + Overhead
)

const (
	infoCBC     = "AES-CBC filename encryption key"
	infoCBCIV   = "AES-CBC filename encryption initialization vector"
	infoChaPloy = "ChaCha20-Poly1305 file content encryption"
)

func cbcKey(key []byte) []byte {
	return deriveKey(key, nil, infoCBC, 32)
}

func cbcIV(key []byte) []byte {
	return deriveKey(key, nil, infoCBCIV, 16)
}

func streamKey(key, nonce []byte) []byte {
	return deriveKey(key, nonce, infoChaPloy, chacha20poly1305.KeySize)
}

func deriveKey(key, nonce []byte, info string, length int) []byte {
	h := hkdf.New(sha256.New, key, nonce, []byte(info))
	k := make([]byte, length)
	if _, err := io.ReadFull(h, k); err != nil {
		panic("crypto: internal error: failed to read from HKDF: " + err.Error())
	}
	return k
}

type aesCbc struct {
	key []byte
	iv  []byte
}

func newAesCbc(key []byte) (*aesCbc, error) {
	cbc := &aesCbc{
		key: cbcKey(key),
		iv:  cbcIV(key),
	}
	_, err := aes.NewCipher(cbc.key)
	if err != nil {
		return nil, err
	}
	return cbc, nil
}

func (cbc *aesCbc) encrypt(p []byte) (c []byte) {
	block, _ := aes.NewCipher(cbc.key)
	mode := cipher.NewCBCEncrypter(block, cbc.iv)
	padded := padPKCS5(p, block.BlockSize())
	c = make([]byte, len(padded))
	mode.CryptBlocks(c, padded)
	return
}

func (cbc *aesCbc) decrypt(c []byte) []byte {
	block, _ := aes.NewCipher(cbc.key)
	mode := cipher.NewCBCDecrypter(block, cbc.iv)
	p := make([]byte, len(c))
	mode.CryptBlocks(p, c)
	return trimPKCS5(p)
}

func padPKCS5(b []byte, bs int) []byte {
	n := bs - len(b)%bs
	pad := bytes.Repeat([]byte{byte(n)}, n)
	return append(b, pad...)
}

func trimPKCS5(b []byte) []byte {
	pad := b[len(b)-1]
	return b[:len(b)-int(pad)]
}