package laosiji

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"errors"
)

var (
	ErrInvalidCipher  = errors.New("laosiji crypto: invalid ciphertext length")
	ErrInvalidPadding = errors.New("laosiji crypto: invalid pkcs7 padding")
)

func AesKey(appkey string) []byte {
	key := make([]byte, 16)
	copy(key, appkey)
	return key
}

func EncryptBase64(plain []byte, appkey string) (string, error) {
	block, err := aes.NewCipher(AesKey(appkey))
	if err != nil {
		return "", err
	}
	padded := pkcs7Pad(plain, block.BlockSize())
	out := make([]byte, len(padded))
	ecbEncrypt(block, out, padded)
	return base64.StdEncoding.EncodeToString(out), nil
}

func DecryptBase64(b64 string, appkey string) ([]byte, error) {
	ct, err := base64.StdEncoding.DecodeString(b64)
	if err != nil {
		return nil, err
	}
	block, err := aes.NewCipher(AesKey(appkey))
	if err != nil {
		return nil, err
	}
	bs := block.BlockSize()
	if len(ct) == 0 || len(ct)%bs != 0 {
		return nil, ErrInvalidCipher
	}
	out := make([]byte, len(ct))
	ecbDecrypt(block, out, ct)
	return pkcs7Unpad(out, bs)
}

func ecbEncrypt(block cipher.Block, dst, src []byte) {
	bs := block.BlockSize()
	for len(src) > 0 {
		block.Encrypt(dst[:bs], src[:bs])
		src = src[bs:]
		dst = dst[bs:]
	}
}

func ecbDecrypt(block cipher.Block, dst, src []byte) {
	bs := block.BlockSize()
	for len(src) > 0 {
		block.Decrypt(dst[:bs], src[:bs])
		src = src[bs:]
		dst = dst[bs:]
	}
}

func pkcs7Pad(data []byte, blockSize int) []byte {
	pad := blockSize - len(data)%blockSize
	return append(data, bytes.Repeat([]byte{byte(pad)}, pad)...)
}

func pkcs7Unpad(data []byte, blockSize int) ([]byte, error) {
	n := len(data)
	if n == 0 || n%blockSize != 0 {
		return nil, ErrInvalidPadding
	}
	pad := int(data[n-1])
	if pad == 0 || pad > blockSize {
		return nil, ErrInvalidPadding
	}
	for i := n - pad; i < n; i++ {
		if int(data[i]) != pad {
			return nil, ErrInvalidPadding
		}
	}
	return data[:n-pad], nil
}