package starcrypto

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"errors"
	"io"
)

const (
	PKCS5PADDING = "PKCS5"
)

func CustomEncryptAesCFB(origData []byte, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	encrypted := make([]byte, aes.BlockSize+len(origData))
	iv := encrypted[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return nil, err
	}
	stream := cipher.NewCFBEncrypter(block, iv)
	stream.XORKeyStream(encrypted[aes.BlockSize:], origData)
	return encrypted, nil
}

func CustomDecryptAesCFB(encrypted []byte, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	if len(encrypted) < aes.BlockSize {
		return nil, errors.New("ciphertext too short")
	}
	iv := encrypted[:aes.BlockSize]
	encrypted = encrypted[aes.BlockSize:]

	stream := cipher.NewCFBDecrypter(block, iv)
	stream.XORKeyStream(encrypted, encrypted)
	return encrypted, nil
}

func CustomEncryptAesCFBNoBlock(origData []byte, key []byte, iv []byte) ([]byte, error) {
	if len(iv) != 16 {
		return nil, errors.New("iv length must be 16")
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	encrypted := make([]byte, len(origData))
	stream := cipher.NewCFBEncrypter(block, iv)
	stream.XORKeyStream(encrypted, origData)
	return encrypted, err
}

func CustomDecryptAesCFBNoBlock(encrypted []byte, key []byte, iv []byte) ([]byte, error) {
	if len(iv) != 16 {
		return nil, errors.New("iv length must be 16")
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	stream := cipher.NewCFBDecrypter(block, iv)
	stream.XORKeyStream(encrypted, encrypted)
	return encrypted, err
}

func EncryptAesCBC(data, key []byte, iv []byte, paddingType string) ([]byte, error) {
	var content []byte
	aesBlockEncrypter, err := aes.NewCipher(key)
	switch paddingType {
	case PKCS5PADDING:
		content = PKCS5Padding(data, aesBlockEncrypter.BlockSize())
	default:
		return nil, errors.New("padding type not supported")
	}
	encrypted := make([]byte, len(content))
	if err != nil {
		return nil, err
	}
	aesEncrypter := cipher.NewCBCEncrypter(aesBlockEncrypter, iv)
	aesEncrypter.CryptBlocks(encrypted, content)
	return encrypted, nil
}

func PKCS5Padding(cipherText []byte, blockSize int) []byte {
	padding := blockSize - len(cipherText)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(cipherText, padText...)
}

func PKCS5Trimming(encrypt []byte) []byte {
	padding := encrypt[len(encrypt)-1]
	if len(encrypt)-int(padding) < 0 {
		return nil
	}
	return encrypt[:len(encrypt)-int(padding)]
}

func DecryptAesCBC(src, key []byte, iv []byte, paddingType string) (data []byte, err error) {
	decrypted := make([]byte, len(src))
	var aesBlockDecrypter cipher.Block
	aesBlockDecrypter, err = aes.NewCipher(key)
	if err != nil {
		println(err.Error())
		return nil, err
	}
	aesDecrypter := cipher.NewCBCDecrypter(aesBlockDecrypter, iv)
	aesDecrypter.CryptBlocks(decrypted, src)
	switch paddingType {
	case PKCS5PADDING:
		return PKCS5Trimming(decrypted), nil
	default:
		return nil, errors.New("padding type not supported")
	}
}