// Package xts implements the XTS cipher mode as specified in IEEE P1619/D16.
//
// XTS mode is typically used for disk encryption, which presents a number of
// novel problems that make more common modes inapplicable. The disk is
// conceptually an array of sectors and we must be able to encrypt and decrypt
// a sector in isolation. However, an attacker must not be able to transpose
// two sectors of plaintext by transposing their ciphertext.
//
// XTS wraps a block cipher with Rogaway's XEX mode in order to build a
// tweakable block cipher. This allows each sector to have a unique tweak and
// effectively create a unique key for each sector.
//
// XTS does not provide any authentication. An attacker can manipulate the
// ciphertext and randomise a block (16 bytes) of the plaintext. This package
// does not implement ciphertext-stealing so sectors must be a multiple of 16
// bytes.
//
// Note that XTS is usually not appropriate for any use besides disk encryption.
// Most users should use an AEAD mode like GCM (from crypto/cipher.NewGCM) instead.

package blockmode

import (
	"xdx.jelly/xgcl/gerrors"
	"xdx.jelly/xgcl/internal/xor"
)

// xts contains an expanded key structure. It is safe for concurrent use if
// the underlying block cipher is safe for concurrent use.
type xts struct {
	k1, k2 EcbBlockMode
	tweak  []byte

	outBuf []byte
}

var _ TernaryCrypter = &xts{}

// xtsBlockSize is the block size that the underlying cipher must have. XTS is
// only defined for 16-byte ciphers.
const xtsBlockSize = 16

// NewXTS 需要使用两个16字节key, 对应入参需要两个EcbEncBlockMode。
func NewXTS(ecb1, ecb2 EcbBlockMode) *xts {
	return &xts{
		k1:     ecb1,
		k2:     ecb2,
		tweak:  make([]byte, 0, xtsBlockSize),
		outBuf: make([]byte, 0, 2*xtsBlockSize),
	}
}

func (x *xts) EncryptInit(tw []byte) error {
	if len(tw) != xtsBlockSize {
		return gerrors.WithAnnotating(ErrXTSEncFailed, "input tweak must be 16 bytes long")
	}
	if cap(x.outBuf) < 2*xtsBlockSize {
		x.outBuf = make([]byte, 0, 2*xtsBlockSize)
	}
	x.outBuf = x.outBuf[:0]
	x.tweak = append(x.tweak[:0], tw...)
	_ = x.k2.EcbEncCryptBlocks(x.tweak, x.tweak)
	return nil
}

func (x *xts) Encrypt(dst []byte, src []byte) ([]byte, error) {
	dst, err := x.EncryptUpdate(dst, src)
	if err != nil {
		return nil, err
	}
	return x.EncryptFinal(dst)
}

func (x *xts) fillTweakMask(mask []byte) {
	for i := 0; i < len(mask); i += xtsBlockSize {
		copy(mask[i:i+xtsBlockSize], x.tweak[:])
		mul2(x.tweak)
	}
}

func (x *xts) EncryptUpdate(dst []byte, in []byte) ([]byte, error) {
	// how many data in hand
	dataLen := len(x.outBuf) + len(in)
	if dataLen < 2*xtsBlockSize {
		x.outBuf = append(x.outBuf, in...)
		return dst, nil
	}
	// how many blocks to process in this update
	m := (dataLen - xtsBlockSize) >> 4

	mask := make([]byte, m<<4)
	x.fillTweakMask(mask)
	ret, out := sliceForAppend(dst, m<<4)
	n := xor.XorBytes(out, mask, x.outBuf)
	// TODO wrong if len(in) = 1
	n = xor.XorBytes(out[n:], mask[n:], in)
	in = in[n:]
	_ = x.k1.EcbEncCryptBlocks(out, out)

	xor.XorBytes(out, out, mask)
	n = 0
	if len(x.outBuf) > m<<4 {
		n = copy(x.outBuf[:xtsBlockSize], x.outBuf[m<<4:])
	}
	x.outBuf = x.outBuf[:n]
	x.outBuf = append(x.outBuf, in...)
	return ret, nil
}

func (x *xts) EncryptFinal(dst []byte) ([]byte, error) {
	if len(x.outBuf) < xtsBlockSize {
		return dst, gerrors.WithAnnotating(ErrXTSEncFailed, "plaintext must be at least 16 bytes")
	}
	in0 := x.outBuf[:xtsBlockSize]
	in1 := x.outBuf[xtsBlockSize:]

	ret, out := sliceForAppend(dst, len(x.outBuf))
	out0 := out[:xtsBlockSize]
	out1 := out[xtsBlockSize:]

	mask := make([]byte, xtsBlockSize)
	x.fillTweakMask(mask)
	xor.XorBytes(out0, mask, in0)
	_ = x.k1.EcbEncCryptBlocks(out0, out0)
	xor.XorBytes(out0, out0, mask)

	if len(in1) > 0 {
		lastBlock := make([]byte, len(in1), xtsBlockSize)
		copy(lastBlock, in1)
		lastBlock = append(lastBlock, out0[len(in1):]...)
		x.fillTweakMask(mask)
		xor.XorBytes(lastBlock, mask, lastBlock)
		_ = x.k1.EcbEncCryptBlocks(lastBlock, lastBlock)
		xor.XorBytes(lastBlock, lastBlock, mask)
		copy(mask, out0)
		copy(out0, lastBlock)
		copy(out1, mask)
	}

	return ret, nil

}
func (x *xts) DecryptInit(tw []byte) error {
	if len(tw) != xtsBlockSize {
		return gerrors.WithAnnotating(ErrXTSDecFailed, "input tweak must be 16 bytes long")
	}
	x.tweak = append(x.tweak[:0], tw...)
	if cap(x.outBuf) < 2*xtsBlockSize {
		x.outBuf = make([]byte, 0, 2*xtsBlockSize)
	}
	x.outBuf = x.outBuf[:0]
	_ = x.k2.EcbEncCryptBlocks(x.tweak, x.tweak)
	return nil
}
func (x *xts) Decrypt(dst []byte, in []byte) ([]byte, error) {
	dst, err := x.DecryptUpdate(dst, in)
	if err != nil {
		return nil, err
	}
	return x.DecryptFinal(dst)
}
func (x *xts) DecryptUpdate(dst []byte, in []byte) ([]byte, error) {
	// how many data in hand
	dataLen := len(x.outBuf) + len(in)
	if dataLen < 2*xtsBlockSize {
		x.outBuf = append(x.outBuf, in...)
		return dst, nil
	}
	// how many blocks to process in this update
	m := (dataLen - xtsBlockSize) >> 4

	mask := make([]byte, m<<4)
	x.fillTweakMask(mask)
	ret, out := sliceForAppend(dst, m<<4)
	n := xor.XorBytes(out, mask, x.outBuf)
	// TODO wrong if len(in) = 1
	n = xor.XorBytes(out[n:], mask[n:], in)
	in = in[n:]
	_ = x.k1.EcbDecCryptBlocks(out, out)

	xor.XorBytes(out, out, mask)
	n = 0
	if len(x.outBuf) > m<<4 {
		n = copy(x.outBuf[:xtsBlockSize], x.outBuf[m<<4:])
	}
	x.outBuf = x.outBuf[:n]
	x.outBuf = append(x.outBuf, in...)
	return ret, nil
}
func (x *xts) DecryptFinal(dst []byte) ([]byte, error) {

	if len(x.outBuf) < xtsBlockSize {
		return dst, gerrors.WithAnnotating(ErrXTSEncFailed, "plaintext must be at least 16 bytes")
	}

	in0 := x.outBuf[:xtsBlockSize]
	in1 := x.outBuf[xtsBlockSize:]

	ret, out := sliceForAppend(dst, len(x.outBuf))
	out0 := out[:xtsBlockSize]
	out1 := out[xtsBlockSize:]

	mask := make([]byte, xtsBlockSize*2)
	x.fillTweakMask(mask)
	mask0 := mask[:xtsBlockSize]
	mask1 := mask[xtsBlockSize:]

	if len(in1) == 0 {
		xor.XorBytes(out0, mask0, in0)
		_ = x.k1.EcbDecCryptBlocks(out0, out0)
		xor.XorBytes(out0, out0, mask0)
	} else {
		xor.XorBytes(out0, mask1, in0)
		_ = x.k1.EcbDecCryptBlocks(out0, out0)
		xor.XorBytes(out0, out0, mask1)

		lastBlock := make([]byte, len(in1), xtsBlockSize)
		copy(lastBlock, in1)
		lastBlock = append(lastBlock, out0[len(in1):]...)
		xor.XorBytes(lastBlock, mask0, lastBlock)
		_ = x.k1.EcbDecCryptBlocks(lastBlock, lastBlock)
		xor.XorBytes(lastBlock, lastBlock, mask0)
		copy(mask, out0)
		copy(out0, lastBlock)
		copy(out1, mask)
	}

	return ret, nil
}

// mul2 multiplies tweak by 2 in GF(2¹²⁸) with an irreducible polynomial of
// x¹²⁸ + x⁷ + x² + x + 1.
func mul2(tweak []byte) {
	var carryIn byte
	for j := range tweak {
		carryOut := tweak[j] >> 7
		tweak[j] = (tweak[j] << 1) + carryIn
		carryIn = carryOut
	}
	if carryIn != 0 {
		// If we have a carry bit then we need to subtract a multiple
		// of the irreducible polynomial (x¹²⁸ + x⁷ + x² + x + 1).
		// By dropping the carry bit, we're subtracting the x^128 term
		// so all that remains is to subtract x⁷ + x² + x + 1.
		// Subtraction (and addition) in this representation is just
		// XOR.
		tweak[0] ^= 1<<7 | 1<<2 | 1<<1 | 1
	}
}