/*
** Copyright (c) 2017 Arnaud Ysmal.  All Rights Reserved.
**
** Redistribution and use in source and binary forms, with or without
** modification, are permitted provided that the following conditions
** are met:
** 1. Redistributions of source code must retain the above copyright
**    notice, this list of conditions and the following disclaimer.
** 2. Redistributions in binary form must reproduce the above copyright
**    notice, this list of conditions and the following disclaimer in the
**    documentation and/or other materials provided with the distribution.
**
** THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS
** OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
** WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
** DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
** FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
** DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
** SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
** HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
** LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
** OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
** SUCH DAMAGE.
 */

package aes

import (
	"crypto/aes"
	"crypto/cipher"

	"github.com/stacktic/cryptodev/cryptodev"
)

const BlockSize = 16

type cryptodevAES struct {
	key []byte
	b   cipher.Block
}

type cryptodevAESCBCEncrypter struct {
	cryptodevAES
	s  *cryptodev.Session
	iv []byte
}

type cryptodevAESCBCDecrypter struct {
	cryptodevAES
	s   *cryptodev.Session
	iv  []byte
	tmp []byte
}

func (c *cryptodevAES) NewCBCDecrypter(iv []byte) cipher.BlockMode {
	var err error
	var s *cryptodev.Session

	if len(iv) != c.BlockSize() {
		panic("cipher.NewCBCDecrypter: IV length must equal block size")
	}

	if s, err = cryptodev.NewSession(cryptodev.CRYPTO_AES_CBC, c.key); err != nil {
		return cipher.NewCBCDecrypter(c.b, iv)
	}

	bm := &cryptodevAESCBCDecrypter{
		cryptodevAES: *c,
		s:            s,
		iv:           make([]byte, BlockSize),
		tmp:          make([]byte, BlockSize),
	}
	copy(bm.iv, iv)
	return bm
}

func (c *cryptodevAES) NewCBCEncrypter(iv []byte) cipher.BlockMode {
	var err error
	var s *cryptodev.Session

	if len(iv) != c.BlockSize() {
		panic("cipher.NewCBCEncrypter: IV length must equal block size")
	}

	if s, err = cryptodev.NewSession(cryptodev.CRYPTO_AES_CBC, c.key); err != nil {
		return cipher.NewCBCEncrypter(c.b, iv)
	}
	bm := &cryptodevAESCBCEncrypter{
		cryptodevAES: *c,
		s:            s,
		iv:           make([]byte, BlockSize),
	}
	copy(bm.iv[:], iv)
	return bm
}

func NewCipher(key []byte) (cipher.Block, error) {
	var err error
	var b cipher.Block

	if b, err = aes.NewCipher(key); err != nil {
		return nil, err
	}
	return &cryptodevAES{key: key, b: b}, nil
}

func (c *cryptodevAES) Encrypt(src []byte, dst []byte) {
	c.b.Encrypt(src, dst)
}

func (c *cryptodevAES) Decrypt(src []byte, dst []byte) {
	c.b.Decrypt(src, dst)
}

func (c *cryptodevAES) BlockSize() int {
	return BlockSize
}

func (c *cryptodevAESCBCDecrypter) CryptBlocks(dst, src []byte) {
	if len(src)%BlockSize != 0 {
		panic("crypto/cipher: input not full blocks")
	}
	if len(dst) < len(src) {
		panic("crypto/cipher: output smaller than input")
	}
	if len(src) > 0 {
		copy(c.tmp, src[len(src)-BlockSize:])
		c.s.Decrypt(dst, src, c.iv)
		c.tmp, c.iv = c.iv, c.tmp
	}
}

func (c *cryptodevAESCBCEncrypter) CryptBlocks(dst, src []byte) {
	if len(src)%BlockSize != 0 {
		panic("crypto/cipher: input not full blocks")
	}
	if len(dst) < len(src) {
		panic("crypto/cipher: output smaller than input")
	}
	if len(src) > 0 {
		c.s.Encrypt(dst, src, c.iv)
		copy(c.iv, dst[len(dst)-BlockSize:])
	}
}