// Copyright (c) Huawei Technologies Co., Ltd. 2020. All rights reserved.
// iSula-Kits licensed under the Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//     http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
// PURPOSE.
// See the Mulan PSL v2 for more details.
// Author: Xiang Li
// Create: 2020-06-02
// Description: This file stores functions which used for aes encrypting and decrypting

package util

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"hash"
	"io"

	"github.com/pkg/errors"
	"golang.org/x/crypto/pbkdf2"
)

const (
	// CryptoKeyLen is secure key length for aes encryption and decryption(AES-256)
	CryptoKeyLen = 32
	// iteration is iteration count to hash
	iteration           = 409600
	aesKeyLenUpperBound = 32
	aesKeyLenLowerBound = 16
)

var (
	errGenCryptoKey = errors.New("generate crypto key failed")
)

// GenerateCryptoKey generates a random key with length s
// if used with AES, the input length can only be 16, 24, 32,
// which stands for AES-128, AES-192, or AES-256.
func GenerateCryptoKey(s int) ([]byte, error) {
	var size int
	if s >= aesKeyLenLowerBound && s <= aesKeyLenUpperBound {
		size = s
	} else {
		size = aesKeyLenLowerBound
	}
	key := make([]byte, size, size)
	if _, err := io.ReadFull(rand.Reader, key); err != nil {
		return nil, errGenCryptoKey
	}

	return key, nil
}

// PBKDF2 is key derivation function to generate one way hash data
// if used with AES, the keyLen can only be 16, 24, 32
// which stands for AES-128, AES-192 or AES-256
// iteration is pre-set to 409600 and salt is generated by random key generator
func PBKDF2(password []byte, keyLen int, h func() hash.Hash) (string, error) {
	if len(password) == 0 {
		return "", errors.New("encrypt empty string failed")
	}
	salt, err := GenerateCryptoKey(CryptoKeyLen)
	if err != nil {
		return "", err
	}

	df := pbkdf2.Key(password, salt, iteration, keyLen, h)

	return hex.EncodeToString(df), nil
}

// EncryptAES encrypts plain text with AES encrypt algorithm(CFB)
func EncryptAES(data string, aeskey string) (string, error) {
	plainText := []byte(data)
	key, err := hex.DecodeString(aeskey)
	if err != nil {
		return "", err
	}

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

	iv, err := GenerateCryptoKey(block.BlockSize())
	if err != nil {
		return "", errors.Errorf("generate rand data for iv failed: %v", err)
	}
	mode := cipher.NewCFBEncrypter(block, iv)
	encryptData := make([]byte, len(plainText), len(plainText))
	mode.XORKeyStream(encryptData, plainText)
	encryptData = append(iv, encryptData...)

	return hex.EncodeToString(encryptData), nil
}

// DecryptAES decrypts text with AES decrypt algorithm(CFB)
func DecryptAES(data string, aeskey string) (string, error) {
	key, err := hex.DecodeString(aeskey)
	if err != nil {
		return "", err
	}

	cipherText, err := hex.DecodeString(data)
	if err != nil {
		return "", err
	}

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

	if len(cipherText) <= block.BlockSize() {
		return "", errors.Errorf("invalid cipher text length %v, it must larger than %v", len(cipherText), block.BlockSize())
	}

	decrypter := cipher.NewCFBDecrypter(block, cipherText[:block.BlockSize()])
	decryptData := make([]byte, len(cipherText)-block.BlockSize(), len(cipherText)-block.BlockSize())
	decrypter.XORKeyStream(decryptData, cipherText[block.BlockSize():])

	return string(decryptData), nil
}
