///
/// Copyright (c) 2018 xdx. All rights reserved.
///
/// \file: encryption.go
///
/// \brief: SM9加解密
///
/// \author: xdx
///

package sm9

import (
	"io"

	"xdx.jelly/xgcl/gerrors"
	"xdx.jelly/xgcl/grand"
	"xdx.jelly/xgcl/sm/sm4"
	"xdx.jelly/xgcl/sm/sm9/errors"
)

type EncType uint32

// 加密类型
const (
	EncTypeKDF    EncType = 0
	EncTypeSM4ECB EncType = 1
	EncTypeSM4CBC EncType = 2
	EncTypeSM4CFB EncType = 4
	EncTypeSM4OFB EncType = 8
	encTypeDummy  EncType = 0xffffffff //无效类型
)

// defaultCipherLength 默认初始分配大小，并不是限定密文长度。
// 一般加密16字节SM4密钥
const defaultCipherLength = 32

// Cipher 密文结构
type Cipher struct {
	EncType EncType             // 加密类型
	IV      [sm4.BlockSize]byte // iv, 当为SM4 CBC,OFB...时有效
	C1      G1
	H       [macSize]byte // C3
	C       []byte        // C2

}

// NewCipher allocate memeory for a cipher
// TODO: NewXXX should also works when replaced by &XXX{}
func NewCipher() *Cipher {
	return &Cipher{
		EncType: encTypeDummy,
		C:       make([]byte, 0, defaultCipherLength),
	}
}

// MarshalBinary return type || X || Y || C3 || len(IV+C2) || IV || C2
func (c *Cipher) MarshalBinary() ([]byte, error) {
	dataLen := 4 + 2*Sm9RefMaxLen + macSize + 4 + len(c.C)
	if (c.EncType == EncTypeSM4CBC || c.EncType == EncTypeSM4CFB || c.EncType == EncTypeSM4OFB) && len(c.IV) == sm4.BlockSize {
		dataLen += sm4.BlockSize
	}
	data := make([]byte, 4, dataLen)
	Endian.PutUint32(data, uint32(c.EncType))
	data = append(data, c.C1.Marshal()...)
	data = append(data, c.H[:]...)
	buf := make([]byte, 4)
	Endian.PutUint32(buf, uint32(len(c.C)))
	data = append(data, buf...)
	if c.EncType == EncTypeSM4CBC && len(c.IV) == sm4.BlockSize {
		data = append(data, c.IV[:]...)
	}
	data = append(data, c.C...)
	return data, nil
}

// UnmarshalBinary implements the encoding.BinaryUnmarshaler interface
func (c *Cipher) UnmarshalBinary(data []byte) error {
	if len(data) < 4+2*Sm9RefMaxLen+macSize+4 {
		return gerrors.WithAnnotating(errors.ErrUnmarshalFailed, "input unmarshal data too short")
	}
	encType := EncType(Endian.Uint32(data))
	data = data[4:]

	var err error
	if data, err = c.C1.Unmarshal(data); err != nil {
		return gerrors.ChainErrors(errors.ErrUnmarshalFailed, err)
	}

	c.EncType = EncType(encType)
	copy(c.H[:], data[:macSize])
	data = data[macSize:]

	// if CBC, clen include the iv
	clen := Endian.Uint32(data)
	data = data[4:]
	// clen too short even for a iv
	if encType == EncTypeSM4CBC || encType == EncTypeSM4CFB || encType == EncTypeSM4OFB {
		if clen < sm4.BlockSize {
			return gerrors.WithAnnotating(errors.ErrUnmarshalFailed, "input unmarshal data of C2 less than a block")
		}
		copy(c.IV[:], data[:sm4.BlockSize])
		data = data[sm4.BlockSize:]
	}

	if len(data) < int(clen) {
		return gerrors.WithAnnotating(errors.ErrUnmarshalFailed, "input unmarshal data too short")
	}
	c.C = append(c.C[:0], data...)
	return nil
}

// Encrypt SM9加密
//
// 输入:
//   - pube: 加密主公钥
//   - encType: 加密类型
//   - optionalIV: 当加密类型为EncTypeSM4CBC、EncTypeSM4CFB、EncTypeSM4OFB时, 输入16字节iv值。也可以输入nil, 则使用rand随机产生。
//   - plain: 加密原文，当为非KDF加密类型时，长度必须是16的倍数，内部不做padding
func Encrypt(encType EncType, pube *MastEncPublicKey, id, plain []byte, rand io.Reader, optionalIV []byte) (*Cipher, error) {
	rnd := make([]byte, Sm9RefMaxLen)
	if rand == nil {
		rand = grand.Reader
	}
	if _, err := rand.Read(rnd); err != nil {
		return nil, err
	}
	if encType != EncTypeKDF && len(plain)%sm4.BlockSize != 0 {
		return nil, gerrors.WithAnnotating(errors.ErrEncFailed, "input plaintext is not multilple of block size, padding first")
	}

	switch encType {
	case EncTypeKDF:
		return encryption(id, plain, pube, nil, rnd, len(plain)+macKeySize, xor, EncTypeKDF), nil
	case EncTypeSM4ECB:
		return encryption(id, plain, pube, nil, rnd, sm4.BlockSize+macKeySize, encEcb, EncTypeSM4ECB), nil
	case EncTypeSM4CBC:
		if len(optionalIV) < numBytes {
			optionalIV = make([]byte, sm4.BlockSize)
			if _, err := rand.Read(optionalIV); err != nil {
				return nil, err
			}
		}
		return encryption(id, plain, pube, optionalIV, rnd, sm4.BlockSize+macKeySize, encCbc, EncTypeSM4CBC), nil
	case EncTypeSM4CFB:
		if len(optionalIV) == 0 {
			optionalIV = make([]byte, sm4.BlockSize)
			if _, err := rand.Read(optionalIV); err != nil {
				return nil, err
			}
		}
		return encryption(id, plain, pube, optionalIV, rnd, sm4.BlockSize+macKeySize, encCfb, EncTypeSM4CFB), nil
	case EncTypeSM4OFB:
		if len(optionalIV) == 0 {
			optionalIV = make([]byte, sm4.BlockSize)
			if _, err := rand.Read(optionalIV); err != nil {
				return nil, err
			}
		}
		return encryption(id, plain, pube, optionalIV, rnd, sm4.BlockSize+macKeySize, encOfb, EncTypeSM4OFB), nil
	default:
		return nil, errors.ErrEncUnsupportedMode
	}
}

// Decrypt SM9解密。分组模式不对解密明文做unpadding
func Decrypt(id []byte, c *Cipher, de *UserEncKey) ([]byte, error) {
	keylen, f, err := DecryptParams(c)
	if err != nil {
		return nil, err
	}
	return decryption(id, c, de, keylen, f)
}

// EncryptionKDF kdf模式加密
//
// Deprecated: use Encrypt/Decrypt instead.
func EncryptionKDF(id, plain []byte, pube *MastEncPublicKey, rnd []byte) (*Cipher, error) {
	return encryption(id, plain, pube, nil, rnd, len(plain)+macKeySize, xor, EncTypeKDF), nil
}

// DecryptionKDF kdf模式加密
//
// Deprecated: use Encrypt/Decrypt instead.
func DecryptionKDF(id []byte, c *Cipher, de *UserEncKey) ([]byte, error) {
	return decryption(id, c, de, len(c.C)+macKeySize, xor)
}

// EncryptionSm4ECB SM4_ECB模式, 不做padding和unpadding
//
// Deprecated: use Encrypt/Decrypt instead.
func EncryptionSm4ECB(id, plain []byte, pube *MastEncPublicKey, rnd []byte) (*Cipher, error) {
	if len(plain)%sm4.BlockSize != 0 {
		return nil, gerrors.WithAnnotating(errors.ErrEncFailed, "input plaintext is not multilple of block size, padding first")
	}
	return encryption(id, plain, pube, nil, rnd, sm4.BlockSize+macKeySize, encEcb, EncTypeSM4ECB), nil
}

// DecryptionSm4ECB SM4_ECB模式, 不做padding和unpadding
//
// Deprecated: use Encrypt/Decrypt instead.
func DecryptionSm4ECB(id []byte, c *Cipher, de *UserEncKey) ([]byte, error) {
	return decryption(id, c, de, sm4.BlockSize+macKeySize, decEcb)
}

// EncryptionSm4CBC SM4_CBC模式，不做padding和unpadding
//
// Deprecated: use Encrypt/Decrypt instead.
func EncryptionSm4CBC(id, plain []byte, pube *MastEncPublicKey, iv, rnd []byte) (*Cipher, error) {
	if len(plain)%sm4.BlockSize != 0 {
		return nil, gerrors.WithAnnotating(errors.ErrEncFailed, "input plaintext is not multilple of block size, padding first")
	}
	return encryption(id, plain, pube, iv, rnd, sm4.BlockSize+macKeySize, encCbc, EncTypeSM4CBC), nil
}

// DecryptionSm4CBC SM4_ECB模式, 不做padding和unpadding
//
// Deprecated: use Encrypt/Decrypt instead.
func DecryptionSm4CBC(id []byte, c *Cipher, de *UserEncKey) ([]byte, error) {
	return decryption(id, c, de, sm4.BlockSize+macKeySize, decCbc)
}

// EncryptionSm4CFB SM4_CFB模式，不做padding和unpadding
//
// Deprecated: use Encrypt/Decrypt instead.
func EncryptionSm4CFB(id, plain []byte, pube *MastEncPublicKey, iv, rnd []byte) (*Cipher, error) {
	if len(plain)%sm4.BlockSize != 0 {
		return nil, gerrors.WithAnnotating(errors.ErrEncFailed, "input plaintext is not multilple of block size, padding first")
	}
	return encryption(id, plain, pube, iv, rnd, sm4.BlockSize+macKeySize, encCfb, EncTypeSM4CFB), nil
}

// DecryptionSm4CFB SM4_CFB模式, 不做padding和unpadding
//
// Deprecated: use Encrypt/Decrypt instead.
func DecryptionSm4CFB(id []byte, c *Cipher, de *UserEncKey) ([]byte, error) {
	return decryption(id, c, de, sm4.BlockSize+macKeySize, decCfb)
}

// EncryptionSm4OFB SM4_OFB模式，不做padding和unpadding
//
// Deprecated: use Encrypt/Decrypt instead.
func EncryptionSm4OFB(id, plain []byte, pube *MastEncPublicKey, iv, rnd []byte) (*Cipher, error) {
	if len(plain)%sm4.BlockSize != 0 {
		return nil, gerrors.WithAnnotating(errors.ErrEncFailed, "input plaintext is not multilple of block size, padding first")
	}
	return encryption(id, plain, pube, iv, rnd, sm4.BlockSize+macKeySize, encOfb, EncTypeSM4OFB), nil
}

// DecryptionSm4OFB SM4_OFB模式, 不做padding和unpadding
//
// Deprecated: use Encrypt/Decrypt instead.
func DecryptionSm4OFB(id []byte, c *Cipher, de *UserEncKey) ([]byte, error) {
	return decryption(id, c, de, sm4.BlockSize+macKeySize, decOfb)
}