package gospake2

import (
	"bytes"
	"fmt"
	"math/big"
	"salsa.debian.org/vasudev/gospake2/ed25519group"
	group "salsa.debian.org/vasudev/gospake2/groups"
	"salsa.debian.org/vasudev/gospake2/integergroup"
	"testing"
)

func TestTranscriptAssymetric(t *testing.T) {
	key := fmt.Sprintf("%x", transcriptAsymmetric(NewPassword("pw"),
		NewIdentityA("idA"), NewIdentityB("idB"),
		[]byte("XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"),
		[]byte("YYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYY"),
		[]byte("KKKKKKKKKKKKKKKKKKKKKKKKKKKKKKKK")))
	expected := "d59d9ba920f7092565cec747b08d5b2e981d553ac32fde0f25e5b4a4cfca3efd"
	if key != expected {
		t.Errorf("key derived is not same as expected: %s\n", key)
	}

}

func TestTranscriptSymmetric(t *testing.T) {
	key := fmt.Sprintf("%x", transcriptSymmetric(NewPassword("pw"), NewIdentityS("idSymmetric"),
		[]byte("XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"),
		[]byte("YYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYY"),
		[]byte("KKKKKKKKKKKKKKKKKKKKKKKKKKKKKKKK")))
	expected := "b0b31e4401aae37d91a9a8bf6fbb1298cafc005ff9142e3ffc5b9799fb11128b"
	if key != expected {
		t.Errorf("key derived from symmetric mode is not as expected: %s\n", key)
	}
}

func TestAsymmetric(t *testing.T) {

	password := NewPassword("password")
	idA := NewIdentityA("idA")
	idB := NewIdentityB("idB")

	sidea := SPAKE2A(password, idA, idB)
	sideb := SPAKE2B(password, idA, idB)

	m1 := sidea.Start()
	m2 := sideb.Start()

	if len(m1) != len(m2) {
		t.Errorf("PAKE message generated by StartA and StartB should be same")
	}

	if m1[0] != 0x41 {
		t.Errorf("Starting byte of PAKE message for A should be 0x41 but found %x\n", m1[0])
	}

	if m2[0] != 0x42 {
		t.Errorf("Starting byte of PAKE message for B should be 0x42 but found %x\n", m2[0])
	}

	if sidea.passwordScalar.Cmp(sideb.passwordScalar) != 0 {
		t.Errorf("Expected passwordScalar for both side A and side B to be same")
	}

	keyA, errA := sidea.Finish(m2)
	if errA != nil {
		t.Errorf("Got error %s while deriving key for side A\n", errA)
	}

	keyB, errB := sideb.Finish(m1)
	if errB != nil {
		t.Errorf("Got error %s while deriving key for side B\n", errB)
	}

	if !bytes.Equal(keyA, keyB) {
		t.Errorf("Expected key derived by both side to be same but they are not\nside A: %x\nside B: %x\n", keyA, keyB)
	}

}

func TestSymmetric(t *testing.T) {

	password := NewPassword("password")
	idS := NewIdentityS("symmetric")
	stateA := SPAKE2Symmetric(password, idS)
	stateB := SPAKE2Symmetric(password, idS)
	m1 := stateA.Start()
	m2 := stateB.Start()

	if len(m1) != len(m2) {
		t.Errorf("PAKE message generated by StartA and StartB should be same")
	}

	if m1[0] != 0x53 {
		t.Errorf("Starting byte of PAKE message for A should be 0x53 but found %x\n", m2[0])
	}

	if m2[0] != 0x53 {
		t.Errorf("Starting byte of PAKE message for B should be 0x53 but found %x\n", m2[0])
	}

	keyA, errA := stateA.Finish(m2)
	if errA != nil {
		t.Errorf("Got error %s while deriving key for side A\n", errA)
	}

	keyB, errB := stateB.Finish(m1)
	if errB != nil {
		t.Errorf("Got error %s while deriving key for side B\n", errB)
	}

	if !bytes.Equal(keyA, keyB) {
		t.Errorf("Expected key derived by both side to be same but they are not\nside A: %x\nside B: %x\n", keyA, keyB)
	}

	stateA = SPAKE2Symmetric(NewPassword("random"), NewIdentityS("idS"))
	stateB = SPAKE2Symmetric(NewPassword("random"), NewIdentityS("idS"))

	m1 = stateA.Start()
	m2 = stateB.Start()

	if len(m1) != len(m2) {
		t.Errorf("PAKE message generated by StartA and StartB should be same")
	}

	if m1[0] != 0x53 {
		t.Errorf("Starting byte of PAKE message for A should be 0x53 but found %x\n", m2[0])
	}

	if m2[0] != 0x53 {
		t.Errorf("Starting byte of PAKE message for B should be 0x53 but found %x\n", m2[0])
	}

	keyA, errA = stateA.Finish(m2)
	if errA != nil {
		t.Errorf("Got error %s while deriving key for side A\n", errA)
	}

	keyB, errB = stateB.Finish(m1)
	if errB != nil {
		t.Errorf("Got error %s while deriving key for side B\n", errB)
	}

	if !bytes.Equal(keyA, keyB) {
		t.Errorf("Expected key derived by both side to be same but they are not\nside A: %x\nside B: %x\n", keyA, keyB)
	}
}

func decimalToScalar(s []byte) *big.Int {
	scalar := new(big.Int)
	scalar.SetBytes(s)
	return scalar.Mod(scalar, ed25519group.L)
}

func TestStartInternal(t *testing.T) {
	x := new(big.Int)
	x.SetString("2611694063369306139794446498317402240796898290761098242657700742213257926693", 0)

	y := new(big.Int)
	y.SetString("7002393159576182977806091886122272758628412261510164356026361256515836884383", 0)

	expectedPwScalar := new(big.Int)
	expectedPwScalar.SetString("3515301705789368674385125653994241092664323519848410154015274772661223168839", 0)

	password := NewPassword("password")
	idA := NewIdentityA("idA")
	idB := NewIdentityB("idB")

	s1 := SPAKE2A(password, idA, idB)
	s1.setXyScalar(x)

	s2 := SPAKE2B(password, idA, idB)
	s2.setXyScalar(y)

	expectedM1 := "416fc960df73c9cf8ed7198b0c9534e2e96a5984bfc5edc023fd24dacf371f2af9"
	expectedM2 := "42354e97b88406922b1df4bea1d7870f17aed3dba7c720b313edae315b00959309"

	m1 := s1.Start()
	if s1.passwordScalar.Cmp(expectedPwScalar) != 0 {
		t.Errorf("Calculated password scalar is not same as expected\n expected: %s\ncalculated: %s\n", expectedPwScalar.Text(10), s1.passwordScalar.Text(10))
	}

	calcM1 := fmt.Sprintf("%x", m1)
	if calcM1 != expectedM1 {
		t.Errorf("Calculated M1 is not same as expected M1\nexpected: %s\n,calculated: %s\n", expectedM1, calcM1)
	}

	m2 := s2.Start()
	if s2.passwordScalar.Cmp(expectedPwScalar) != 0 {
		t.Errorf("Calculated password scalar is not same as expected\n expected: %s\ncalculated: %s\n", expectedPwScalar.Text(10), s2.passwordScalar.Text(10))
	}

	calcM2 := fmt.Sprintf("%x", m2)
	if calcM2 != expectedM2 {
		t.Errorf("Calculated M2 is not same as expected M2\nexpected: %s\n,calculated: %s\n", expectedM2, calcM2)
	}

	key1, err1 := s1.Finish(m2)
	key2, err2 := s2.Finish(m1)

	if err1 != nil {
		t.Errorf("Encountered error while deriving key from side 1: %s\n", err1)
	}

	if err2 != nil {
		t.Errorf("Encountered error while deriving key from side 1: %s\n", err2)
	}

	if !bytes.Equal(key1, key2) {
		t.Errorf("Could not reach to same key from both side\n")
		t.Errorf("Side 1: %x\n", key1)
		t.Errorf("Side 2: %x\n", key2)
	}

}

func TestErrorConditions(t *testing.T) {
	s := SPAKE2A(NewPassword("password"), NewIdentityA("alice"), NewIdentityB("bob"))

	m := s.Start()
	_, err := s.Finish(m)
	if err == nil {
		t.Errorf("Should have detected reflection attack!")
	}

	switch err.(type) {
	case *SPAKEErr:
		// do nothing
	default:
		t.Errorf("Was expecting SPAKEError type")
	}

	reflecterr := err.(*SPAKEErr)
	if reflecterr.Kind() != ReflectionAttempt {
		t.Errorf("Error should be reflection type but found: %d\n %s\n", reflecterr.Kind(), err)
	}

	s1 := SPAKE2A(NewPassword("password"), NewIdentityA("alice"), NewIdentityB("bob"))
	s2 := SPAKE2B(NewPassword("password"), NewIdentityA("alice"), NewIdentityB("bob"))

	m1 := s1.Start()
	m2 := s2.Start()

	m1[0] = byte(sideB)
	m2[0] = byte(sideA)
	_, berr1 := s1.Finish(m2)
	if berr1 == nil {
		t.Errorf("Was expecting BadSide error")
	}

	b1 := berr1.(*SPAKEErr)
	if b1.Kind() != BadSide {
		t.Errorf("Was expecting BadSide error but got %d\n%s\n", b1.Kind(), berr1)
	}

	_, berr2 := s2.Finish(m1)
	if berr2 == nil {
		t.Errorf("Was expecting BadSide error")
	}

	b2 := berr2.(*SPAKEErr)
	if b2.Kind() != BadSide {
		t.Errorf("Was expecting BadSide error but got %d\n%s\n", b2.Kind(), berr2)
	}

	s = SPAKE2Symmetric(NewPassword("password"), NewIdentityS("idSymmetric"))
	m = s.Start()
	_, err = s.Finish(m1)

	if err == nil {
		t.Errorf("Was expecting BadSide error")
	}

	serr := err.(*SPAKEErr)
	if serr.Kind() != BadSide {
		t.Errorf("Was expecting BadSide error but got %d\n%s\n", b2.Kind(), berr2)
	}

	_, err = s.Finish(m2)

	if err == nil {
		t.Errorf("Was expecting BadSide error")
	}

	serr = err.(*SPAKEErr)
	if serr.Kind() != BadSide {
		t.Errorf("Was expecting BadSide error but got %d\n%s\n", b2.Kind(), berr2)
	}

}

func TestIntegerGroupSPAKE2(t *testing.T) {
	groups := []integergroup.IntegerGroup{integergroup.I1024,
		integergroup.I2048, integergroup.I3072}
	password := NewPassword("password")
	alice := NewIdentityA("Alice")
	bob := NewIdentityB("bob")
	for _, grp := range groups {
		sym1 := SPAKE2A(password, alice, bob)
		sym2 := SPAKE2B(password, alice, bob)

		sym1.SetGroup(grp)
		sym2.SetGroup(grp)

		m1 := sym1.Start()
		m2 := sym2.Start()

		k1, err1 := sym1.Finish(m2)
		k2, err2 := sym2.Finish(m1)

		if err1 != nil {
			t.Errorf("Failed to derive key for side A: %s", err1)
		}

		if err2 != nil {
			t.Errorf("Failed to derive key for side A: %s", err2)
		}

		if !bytes.Equal(k1, k2) {
			t.Errorf("Failed to derive the session key\n")
			t.Errorf("k1: %x\n", k1)
			t.Errorf("k2: %x\n", k2)
		}
	}
}

func testSpake2Asymmetric(grp group.Group) {
	password := NewPassword("SomeRandomP455w0rd")
	idA := NewIdentityA("Alice")
	idB := NewIdentityB("Bob")

	s1 := SPAKE2A(password, idA, idB)
	s2 := SPAKE2B(password, idA, idB)

	switch grp.(type) {
	case integergroup.IntegerGroup:
		s1.SetGroup(grp)
		s2.SetGroup(grp)
	}
	m1 := s1.Start()
	m2 := s2.Start()

	_, _ = s1.Finish(m2)
	_, _ = s2.Finish(m1)

}

func testSpake2Symmetric(grp group.Group) {
	password := NewPassword("SomeRandomP455w0rd")
	id := NewIdentityS("symmetric")

	s1 := SPAKE2Symmetric(password, id)
	s2 := SPAKE2Symmetric(password, id)

	switch grp.(type) {
	case integergroup.IntegerGroup:
		s1.SetGroup(grp)
		s2.SetGroup(grp)
	}

	m1 := s1.Start()
	m2 := s2.Start()
	_, _ = s1.Finish(m2)
	_, _ = s2.Finish(m1)
}

func BenchmarkSPAKE2Ed25519Asymmetric(b *testing.B) {
	for i := 0; i < b.N; i++ {
		testSpake2Asymmetric(ed25519group.Ed25519{})
	}
}

func BenchmarkSPAKE21024Asymmetric(b *testing.B) {
	for i := 0; i < b.N; i++ {
		testSpake2Asymmetric(integergroup.I1024)
	}
}

func BenchmarkSPAKE22048Asymmetric(b *testing.B) {
	for i := 0; i < b.N; i++ {
		testSpake2Asymmetric(integergroup.I2048)
	}
}

func BenchmarkSPAKE23072Asymmetric(b *testing.B) {
	for i := 0; i < b.N; i++ {
		testSpake2Asymmetric(integergroup.I3072)
	}
}

func BenchmarkSPAKE2Ed25519Symmetric(b *testing.B) {
	for i := 0; i < b.N; i++ {
		testSpake2Symmetric(ed25519group.Ed25519{})
	}
}

func BenchmarkSPAKE21024Symmetric(b *testing.B) {
	for i := 0; i < b.N; i++ {
		testSpake2Symmetric(integergroup.I1024)
	}
}

func BenchmarkSPAKE22048Symmetric(b *testing.B) {
	for i := 0; i < b.N; i++ {
		testSpake2Symmetric(integergroup.I2048)
	}
}

func BenchmarkSPAKE23072Symmetric(b *testing.B) {
	for i := 0; i < b.N; i++ {
		testSpake2Symmetric(integergroup.I3072)
	}
}