package cryptstate

import (
	"bytes"
	"crypto/aes"
	"testing"
)

func TestEncrypt(t *testing.T) {
	msg := [15]byte{
		0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
	}
	key := [aes.BlockSize]byte{
		0x96, 0x8b, 0x1b, 0x0c, 0x53, 0x1e, 0x1f, 0x80, 0xa6, 0x1d, 0xcb, 0x27, 0x94, 0x09, 0x6f, 0x32,
	}
	eiv := [aes.BlockSize]byte{
		0x1e, 0x2a, 0x9b, 0xd0, 0x2d, 0xa6, 0x8e, 0x46, 0x26, 0x85, 0x83, 0xe9, 0x14, 0x2a, 0xff, 0x2a,
	}
	div := [aes.BlockSize]byte{
		0x73, 0x99, 0x9d, 0xa2, 0x03, 0x70, 0x00, 0x96, 0xef, 0x55, 0x06, 0x7a, 0x8b, 0xbe, 0x00, 0x07,
	}
	expected := [19]byte{
		0x1f, 0xfc, 0xdd, 0xb4, 0x68, 0x13, 0x68, 0xb7, 0x92, 0x67, 0xca, 0x2d, 0xba, 0xb7, 0x0d, 0x44, 0xdf, 0x32, 0xd4,
	}
	expected_eiv := [aes.BlockSize]byte{
		0x1f, 0x2a, 0x9b, 0xd0, 0x2d, 0xa6, 0x8e, 0x46, 0x26, 0x85, 0x83, 0xe9, 0x14, 0x2a, 0xff, 0x2a,
	}

	cs, err := New()
	if err != nil {
		t.Errorf("%v", err)
	}

	out := make([]byte, 19)
	cs.SetKey(key[0:], eiv[0:], div[0:])
	cs.Encrypt(out[0:], msg[0:])

	if !bytes.Equal(out[0:], expected[0:]) {
		t.Errorf("Mismatch in output")
	}

	if !bytes.Equal(cs.EncryptIV[0:], expected_eiv[0:]) {
		t.Errorf("EIV mismatch")
	}
}

func TestDecrypt(t *testing.T) {
	key := [aes.BlockSize]byte{
		0x96, 0x8b, 0x1b, 0x0c, 0x53, 0x1e, 0x1f, 0x80, 0xa6, 0x1d, 0xcb, 0x27, 0x94, 0x09, 0x6f, 0x32,
	}
	eiv := [aes.BlockSize]byte{
		0x1e, 0x2a, 0x9b, 0xd0, 0x2d, 0xa6, 0x8e, 0x46, 0x26, 0x85, 0x83, 0xe9, 0x14, 0x2a, 0xff, 0x2a,
	}
	div := [aes.BlockSize]byte{
		0x73, 0x99, 0x9d, 0xa2, 0x03, 0x70, 0x00, 0x96, 0xef, 0x55, 0x06, 0x7a, 0x8b, 0xbe, 0x00, 0x07,
	}
	crypted := [19]byte{
		0x1f, 0xfc, 0xdd, 0xb4, 0x68, 0x13, 0x68, 0xb7, 0x92, 0x67, 0xca, 0x2d, 0xba, 0xb7, 0x0d, 0x44, 0xdf, 0x32, 0xd4,
	}
	expected := [15]byte{
		0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
	}
	post_div := [aes.BlockSize]byte{
		0x1f, 0x2a, 0x9b, 0xd0, 0x2d, 0xa6, 0x8e, 0x46, 0x26, 0x85, 0x83, 0xe9, 0x14, 0x2a, 0xff, 0x2a,
	}

	cs, err := New()
	if err != nil {
		t.Errorf("%v", err)
	}

	out := make([]byte, 15)
	cs.SetKey(key[0:], div[0:], eiv[0:])
	cs.Decrypt(out[0:], crypted[0:])

	if !bytes.Equal(out[0:], expected[0:]) {
		t.Errorf("Mismatch in output")
	}

	if !bytes.Equal(cs.DecryptIV[0:], post_div[0:]) {
		t.Errorf("Mismatch in DIV")
	}
}