package codec

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha1"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"io"

	"git.wishpal.cn/wishpal_ironfan/xframe/component/logger"
)

// CfbDecrypt decrypts src with the given key.
func CfbDecrypt(cipherText string, key string) (plainText string, err error) {
	if len(cipherText) == 0 || len(cipherText) < 6 || cipherText[:6] != "crypt-" {
		err = errors.New("illegal ciphertext")
		return
	}
	cipherText = string(cipherText[6:])
	var block cipher.Block
	keyBytes := hashBytes(key)
	cipherTextBytes, _ := hex.DecodeString(cipherText)
	block, err = aes.NewCipher(keyBytes)
	if err != nil {
		return
	}

	if len(cipherTextBytes) < aes.BlockSize {
		err = errors.New("ciphertext too short")
		return
	}

	iv := cipherTextBytes[:aes.BlockSize]
	cipherTextBytes = cipherTextBytes[aes.BlockSize:]
	stream := cipher.NewCFBDecrypter(block, iv)

	plainTextBytes := make([]byte, len(cipherTextBytes))
	stream.XORKeyStream(plainTextBytes, cipherTextBytes)
	plainText = string(plainTextBytes)
	return
}

// CbcEncrypt encrypts dat with the given key.
func CbcEncrypt(key []byte, dat []byte, fixed bool) ([]byte, error) {
	res := cbcPadding(dat, fixed)
	iv := res[:aes.BlockSize]
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	buf := res[aes.BlockSize:]
	cipher.NewCBCEncrypter(block, iv).CryptBlocks(buf, buf)
	return res, nil
}

// CbcDecrypt decrypts dat with the given key.
func CbcDecrypt(key, dat []byte) ([]byte, error) {
	if len(dat) < aes.BlockSize*2 {
		return nil, fmt.Errorf("invalid dat")
	}

	iv := dat[:aes.BlockSize]
	padNum := int(iv[1])
	if padNum <= 0 || padNum > aes.BlockSize {
		return nil, fmt.Errorf("invalid dat")
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	res := make([]byte, len(dat)-aes.BlockSize)
	decrypt := cipher.NewCBCDecrypter(block, iv)
	if len(dat[aes.BlockSize:])%decrypt.BlockSize() != 0 {
		logger.Warnf("decrypt_post_data, key: %v, data: %v", base64.StdEncoding.EncodeToString(key),
			base64.StdEncoding.EncodeToString(dat))
		return nil, fmt.Errorf("crypto/cipher: input not full blocks")
	}
	decrypt.CryptBlocks(res, dat[aes.BlockSize:])
	return res[padNum:], nil
}

func cbcPadding(dat []byte, fixed bool) []byte {
	padNum := aes.BlockSize - (len(dat) % aes.BlockSize)

	//iv+padding+dat
	res := make([]byte, aes.BlockSize+padNum+len(dat))
	iv := res[:aes.BlockSize]
	if !fixed {
		io.ReadFull(rand.Reader, iv)
	}
	iv[1] = byte(padNum)
	padding := res[aes.BlockSize : aes.BlockSize+padNum]
	if !fixed {
		io.ReadFull(rand.Reader, padding)
	}
	
	copy(res[aes.BlockSize+padNum:], dat)
	return res
}

func hashBytes(key string) (hash []byte) {
	h := sha1.New()
	io.WriteString(h, key)
	hashStr := hex.EncodeToString(h.Sum(nil))
	hash = []byte(hashStr)[:32]
	return
}