// Copyright 2018 The go-ethereum Authors
// This file is part of the go-ethereum library.
//
// The go-ethereum library is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// The go-ethereum library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.

package encryption

import (
	"bytes"
	"crypto/rand"
	"testing"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto/sha3"
)

var expectedTransformedHex = "[encoded data, 8192 characters]"

var hashFunc = sha3.NewKeccak256

func TestEncryptDataLongerThanPadding(t *testing.T) {
	enc := New(4095, uint32(0), hashFunc)

	data := make([]byte, 4096)
	key := make([]byte, 32)

	expectedError := "Data length longer than padding, data length 4096 padding 4095"

	_, err := enc.Encrypt(data, key)
	if err == nil || err.Error() != expectedError {
		t.Fatalf("Expected error \"%v\" got \"%v\"", expectedError, err)
	}
}

func TestEncryptDataZeroPadding(t *testing.T) {
	enc := New(0, uint32(0), hashFunc)

	data := make([]byte, 2048)
	key := make([]byte, 32)

	encrypted, err := enc.Encrypt(data, key)
	if err != nil {
		t.Fatalf("Expected no error got %v", err)
	}
	if len(encrypted) != 2048 {
		t.Fatalf("Encrypted data length expected \"%v\" got %v", 2048, len(encrypted))
	}
}

func TestEncryptDataLengthEqualsPadding(t *testing.T) {
	enc := New(4096, uint32(0), hashFunc)

	data := make([]byte, 4096)
	key := make([]byte, 32)

	encrypted, err := enc.Encrypt(data, key)
	if err != nil {
		t.Fatalf("Expected no error got %v", err)
	}
	encryptedHex := common.Bytes2Hex(encrypted)
	expectedTransformed := common.Hex2Bytes(expectedTransformedHex)

	if !bytes.Equal(encrypted, expectedTransformed) {
		t.Fatalf("Expected %v got %v", expectedTransformedHex, encryptedHex)
	}
}

func TestEncryptDataLengthSmallerThanPadding(t *testing.T) {
	enc := New(4096, uint32(0), hashFunc)

	data := make([]byte, 4080)
	key := make([]byte, 32)

	encrypted, err := enc.Encrypt(data, key)
	if err != nil {
		t.Fatalf("Expected no error got %v", err)
	}
	if len(encrypted) != 4096 {
		t.Fatalf("Encrypted data length expected %v got %v", 4096, len(encrypted))
	}
}

func TestEncryptDataCounterNonZero(t *testing.T) {
	// TODO
}

func TestDecryptDataLengthNotEqualsPadding(t *testing.T) {
	enc := New(4096, uint32(0), hashFunc)

	data := make([]byte, 4097)
	key := make([]byte, 32)

	expectedError := "Data length different than padding, data length 4097 padding 4096"

	_, err := enc.Decrypt(data, key)
	if err == nil || err.Error() != expectedError {
		t.Fatalf("Expected error \"%v\" got \"%v\"", expectedError, err)
	}
}

func TestEncryptDecryptIsIdentity(t *testing.T) {
	testEncryptDecryptIsIdentity(t, 2048, 0, 2048, 32)
	testEncryptDecryptIsIdentity(t, 4096, 0, 4096, 32)
	testEncryptDecryptIsIdentity(t, 4096, 0, 1000, 32)
	testEncryptDecryptIsIdentity(t, 32, 10, 32, 32)
}

func testEncryptDecryptIsIdentity(t *testing.T, padding int, initCtr uint32, dataLength int, keyLength int) {
	enc := New(padding, initCtr, hashFunc)

	data := make([]byte, dataLength)
	rand.Read(data)

	key := make([]byte, keyLength)
	rand.Read(key)

	encrypted, err := enc.Encrypt(data, key)
	if err != nil {
		t.Fatalf("Expected no error got %v", err)
	}

	decrypted, err := enc.Decrypt(encrypted, key)
	if err != nil {
		t.Fatalf("Expected no error got %v", err)
	}
	if len(decrypted) != padding {
		t.Fatalf("Expected decrypted data length %v got %v", padding, len(decrypted))
	}

	// we have to remove the extra bytes which were randomly added to fill until padding
	if len(data) < padding {
		decrypted = decrypted[:len(data)]
	}

	if !bytes.Equal(data, decrypted) {
		t.Fatalf("Expected decrypted %v got %v", common.Bytes2Hex(data), common.Bytes2Hex(decrypted))
	}
}
