package cbcmac

import (
	"bytes"
	"crypto/cipher"
	"errors"

	"xdx.jelly/xgcl/utils/padding"
)

// CalculateMAC 计算输入iv和data的cbc Mac.
// 备注:
// 1. data必须是block.BlockSize的整数倍
// 2. iv必须是block.BlockSize大小
// 3. 第一个分组，iv输入全0
// 4. 第i个分组调用后，iv值被更新, 作为第i+1次分组的输入iv
// 5. mac必须为block.BlockSize大小
func CalculateMAC(data []byte, iv []byte, block cipher.Block, mac []byte) error {
	blockSize := block.BlockSize()
	if len(data) == 0 || len(data)%blockSize != 0 || len(iv) != blockSize || len(mac) != blockSize {
		return errors.New("input arguments invalid")
	}
	blockMode := cipher.NewCBCEncrypter(block, iv)

	blockMode.CryptBlocks(mac, data[:blockSize])
	data = data[blockSize:]

	for len(data) > 0 {
		blockMode.CryptBlocks(mac, data[:blockSize])
		data = data[blockSize:]
	}
	return nil
}

const maxChunk = 64

type CbcMac struct {
	mac       []byte
	pad       padding.Padding
	block     cipher.Block
	blockMode cipher.BlockMode
	x         [maxChunk]byte
	nx        int
}

// pad 填充方式。忽略的话为不填充
func New(block cipher.Block, pad ...padding.Padding) (*CbcMac, error) {
	var p padding.Padding
	if len(pad) > 0 {
		p = pad[0]
	} else {
		p = padding.NonePadding{}
	}
	return &CbcMac{
		pad:       p,
		block:     block,
		blockMode: cipher.NewCBCEncrypter(block, make([]byte, block.BlockSize())),
		mac:       make([]byte, block.BlockSize()),
	}, nil
}

// ComputeMAC computes message authentication code (MAC) for the given data.
// data must paded
func (c *CbcMac) ComputeMAC(data []byte) ([]byte, error) {
	c.Reset()
	_, err := c.Write(data)
	if err != nil {
		return nil, err
	}
	defer c.Reset()
	return c.Mac([]byte{})
}

// VerifyMAC verifies whether the given MAC is a correct message authentication
// code (MAC) the given data.
func (c *CbcMac) VerifyMAC(mac []byte, data []byte) error {
	c.Reset()

	mac2, err := c.ComputeMAC(data)
	if err != nil {
		return err
	}
	if bytes.Compare(mac2, mac) != 0 {
		return errors.New("Mac verify failed")
	}
	return nil
}

func (c *CbcMac) Write(p []byte) (n int, err error) {
	blockSize := c.block.BlockSize()

	if len(p)+c.nx < blockSize {
		copy(c.x[c.nx:], p)
		c.nx += len(p)
	} else {
		n = copy(c.x[c.nx:blockSize], p)
		p = p[n:]
		c.blockMode.CryptBlocks(c.mac, c.x[:blockSize])

		for len(p) >= blockSize {
			c.blockMode.CryptBlocks(c.mac, p[:blockSize])
			p = p[blockSize:]
		}
		c.nx = copy(c.x[:], p)
	}
	return len(p), nil
}

// Mac appends the current mac to b and returns the resulting slice.
// tail is the un-computed data c.x[:c.nx].
func (c *CbcMac) Mac(b []byte) ([]byte, error) {
	x := c.pad.Pad(c.x[:c.nx], c.block.BlockSize())
	if len(x)%c.block.BlockSize() != 0 {
		return nil, errors.New("invalid data length with no padding")
	}
	c.blockMode.CryptBlocks(c.mac, x)
	return append(b, c.mac...), nil
}

// Sum is used to compute the mac but return nil while an error occur.
// It's prefered to use Mac instead.
func (c *CbcMac) Sum(b []byte) []byte {
	m, err := c.Mac(b)
	if err != nil {
		return nil
	}
	return m
}

// Reset resets the Hash to its initial state.
func (c *CbcMac) Reset() {
	c.mac = make([]byte, c.block.BlockSize())
	c.blockMode = cipher.NewCBCEncrypter(c.block, c.mac)
}

// Size returns the number of bytes Sum will return.
func (c *CbcMac) Size() int {
	return c.block.BlockSize()
}

// BlockSize returns the hash's underlying block size.
// The Write method must be able to accept any amount
// of data, but it may operate more efficiently if all writes
// are a multiple of the block size.
func (c *CbcMac) BlockSize() int {
	return c.block.BlockSize()
}