package encrypt

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

// Secret is the string to be encrypted
type Secret string

func (s Secret) String() string {
	return string(s)
}

// Encryptor is a basic encryption interface
type Encryptor interface {
	Encrypt(secret Secret) ([]byte, error)
	Decrypt(cyphertext []byte) (Secret, error)
}

// EncryptBase64 encryptes plain text and then encodes it as base64
func EncryptBase64(enc Encryptor, text Secret) (string, error) {
	cipherText, err := enc.Encrypt(text)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(cipherText), nil
}

// DecryptBase64 decrypt base64 encode ciphertext
func DecryptBase64(enc Encryptor, encryptedEncoded string) (Secret, error) {
	encryptedEncodedBytes := []byte(encryptedEncoded)
	encryptedBytes := make([]byte, base64.URLEncoding.DecodedLen(len(encryptedEncodedBytes)))
	bytesWritten, err := base64.StdEncoding.Decode(encryptedBytes, encryptedEncodedBytes)
	if err != nil {
		return "", err
	}
	rootPass, err := enc.Decrypt(encryptedBytes[:bytesWritten])
	if err != nil {
		return "", err
	}
	return Secret(string(rootPass)), nil
}

type AESCrypt struct {
	Block      cipher.Block
	CipherText []byte
}

func NewAESCrypt(cipherText []byte) (*AESCrypt, error) {
	block, err := aes.NewCipher(cipherText)
	if err != nil {
		return nil, err
	}

	return &AESCrypt{
		Block:      block,
		CipherText: cipherText,
	}, nil
}

func (s *AESCrypt) Encrypt(text []byte) []byte {
	return s.EncryptWithIv(text, s.CipherText[:s.Block.BlockSize()])
}

// https://gist.github.com/yingray/57fdc3264b1927ef0f984b533d63abab
func (s *AESCrypt) EncryptWithIv(text []byte, iv []byte) []byte {
	text = s.pkcs5Padding(text, s.Block.BlockSize())
	crypted := make([]byte, len(text))
	encryptor := cipher.NewCBCEncrypter(s.Block, iv)
	encryptor.CryptBlocks(crypted, text)
	return crypted
}

func (s *AESCrypt) Decrypt(crypted []byte) []byte {
	return s.DecryptWithIv(crypted, s.CipherText[:s.Block.BlockSize()])
}

func (s *AESCrypt) DecryptWithIv(crypted []byte, iv []byte) []byte {
	text := make([]byte, len(crypted))
	decryptor := cipher.NewCBCDecrypter(s.Block, iv)
	decryptor.CryptBlocks(text, crypted)
	return s.pkcs5UnPadding(text)
}

func (s *AESCrypt) pkcs5Padding(text []byte, blockSize int) []byte {
	padding := blockSize - len(text)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(text, padtext...)
}

func (s *AESCrypt) pkcs5UnPadding(cryptText []byte) []byte {
	length := len(cryptText)
	unpadding := int(cryptText[length-1])
	return cryptText[:(length - unpadding)]
}