package node

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

	"golang.org/x/crypto/nacl/box"
)

func newControlCipherForTesting() (c1, c2 *controlCipher) {
	pubKey1, privKey1, err := box.GenerateKey(rand.Reader)
	if err != nil {
		panic(err)
	}

	pubKey2, privKey2, err := box.GenerateKey(rand.Reader)
	if err != nil {
		panic(err)
	}

	return newControlCipher(privKey1[:], pubKey2[:]),
		newControlCipher(privKey2[:], pubKey1[:])
}

func TestControlCipher(t *testing.T) {
	c1, c2 := newControlCipherForTesting()

	maxSizePlaintext := make([]byte, bufferSize-controlHeaderSize-controlCipherOverhead)
	rand.Read(maxSizePlaintext)

	testCases := [][]byte{
		make([]byte, 0),
		{1},
		{255},
		{1, 2, 3, 4, 5},
		[]byte("Hello world"),
		maxSizePlaintext,
	}

	for _, plaintext := range testCases {
		h1 := header{
			StreamID: controlStreamID,
			Counter:  235153,
			SourceIP: 4,
			DestIP:   88,
		}

		encrypted := make([]byte, bufferSize)

		encrypted = c1.Encrypt(h1, plaintext, encrypted)

		h2 := header{}
		h2.Parse(encrypted)
		if !reflect.DeepEqual(h1, h2) {
			t.Fatal(h1, h2)
		}

		decrypted, ok := c2.Decrypt(encrypted, make([]byte, bufferSize))
		if !ok {
			t.Fatal(ok)
		}

		if !bytes.Equal(decrypted, plaintext) {
			t.Fatal("not equal")
		}
	}
}

func TestControlCipher_ShortCiphertext(t *testing.T) {
	c1, _ := newControlCipherForTesting()
	shortText := make([]byte, controlHeaderSize+controlCipherOverhead-1)
	rand.Read(shortText)
	_, ok := c1.Decrypt(shortText, make([]byte, bufferSize))
	if ok {
		t.Fatal(ok)
	}
}

func BenchmarkControlCipher_Encrypt(b *testing.B) {
	c1, _ := newControlCipherForTesting()
	h1 := header{
		Counter:  235153,
		SourceIP: 4,
		DestIP:   88,
	}

	plaintext := make([]byte, bufferSize-controlHeaderSize-controlCipherOverhead)
	rand.Read(plaintext)

	encrypted := make([]byte, bufferSize)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		encrypted = c1.Encrypt(h1, plaintext, encrypted)
	}
}

func BenchmarkControlCipher_Decrypt(b *testing.B) {
	c1, c2 := newControlCipherForTesting()

	h1 := header{
		Counter:  235153,
		SourceIP: 4,
		DestIP:   88,
	}

	plaintext := make([]byte, bufferSize-controlHeaderSize-controlCipherOverhead)
	rand.Read(plaintext)

	encrypted := make([]byte, bufferSize)

	encrypted = c1.Encrypt(h1, plaintext, encrypted)

	decrypted := make([]byte, bufferSize)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		decrypted, _ = c2.Decrypt(encrypted, decrypted)
	}
}