package main

import (
	"bytes"
	"encoding/hex"
	"encoding/json"
	"io/ioutil"

	"github.com/henrydcase/nobs/dh/csidh"
	"github.com/henrydcase/sidh_torture/csidh/ref/go-wrapper"
)

// Possible values for "Status"
const (
	Valid               = iota // Indicates that shared secret must be agreed correctly
	ValidPublicKey2            // Public key 2 must succeed validation
	InvalidSharedSecret        // Calculated shared secret must be different than test vector
	InvalidPublicKey1          // Public key 1 generated from private key must be different than test vector
	InvalidPublicKey2          // Public key 2 must fail validation
)

var StatusValues = map[int]string{
	Valid:               "valid",
	ValidPublicKey2:     "valid_public_key2",
	InvalidSharedSecret: "invalid_shared_secret",
	InvalidPublicKey1:   "invalid_public_key1",
	InvalidPublicKey2:   "invalid_public_key2",
}

type TestVector struct {
	Id      int    `json:"Id"`
	Pk1     string `json:"Pk1"`
	Pr1     string `json:"Pr1"`
	Pk2     string `json:"Pk2"`
	Ss      string `json:"Ss"`
	Status  string `json:"status"`
	Comment string `json:"comment"`
}

type TestVectors struct {
	Vectors []TestVector `json:"Vectors"`
}

// R is a reference to C implementation.
var R wrapper.Ref

// createValid creates 'num' of TestVector's. Each vector contains
// valid data.
func createValid(num int) TestVector {
	var tv TestVector
	var ss [csidh.SharedSecretSize]byte
	var pk [csidh.PublicKeySize]byte
	var pr [csidh.PrivateKeySize]byte

	prA := R.KeygenPrv()
	pkA := R.KeygenPub(&prA)
	prB := R.KeygenPrv()
	pkB := R.KeygenPub(&prB)

	R.Derive(ss[:], &pkB, &prA)

	tv.Id = num
	tv.Status = StatusValues[Valid]

	tv.Ss = hex.EncodeToString(ss[:])

	prA.Export(pr[:])
	tv.Pr1 = hex.EncodeToString(pr[:])

	pkA.Export(pk[:])
	tv.Pk1 = hex.EncodeToString(pk[:])

	for i, _ := range pk {
		pk[i] = 0
	}

	pkB.Export(pk[:])
	tv.Pk2 = hex.EncodeToString(pk[:])

	return tv
}

func createNegativeSharedSecret(vectors *TestVectors) {
	n := len(vectors.Vectors)
	tv := createValid(n)
	ss, err := hex.DecodeString(tv.Ss)
	if err != nil {
		panic("Can't decode shared secret")
	}

	for i := 0; i < csidh.SharedSecretSize; i++ {
		var newSs [csidh.SharedSecretSize]byte
		copy(newSs[:], ss[:])
		newSs[i] = ss[i] ^ ss[(i+1)%csidh.SharedSecretSize]

		if bytes.Equal(newSs[:], ss) {
			tv.Status = StatusValues[Valid]
		} else {
			tv.Status = StatusValues[InvalidSharedSecret]
		}
		tv.Ss = hex.EncodeToString(newSs[:])
		tv.Id = n + i
		vectors.Vectors = append(vectors.Vectors, tv)
	}
}

// Public key validation fails
func createNegativePk2(vectors *TestVectors) {
	n := len(vectors.Vectors)
	tv := createValid(n)
	pk, err := hex.DecodeString(tv.Pk2)
	if err != nil {
		panic("Can't decode public key 2")
	}

	for i := 0; i < csidh.PublicKeySize; i++ {
		var newPk [csidh.PublicKeySize]byte

		// Modify good public key so it's probably no longer valid
		copy(newPk[:], pk[:])
		newPk[i] = pk[i] ^ pk[(i+1)%csidh.PublicKeySize]

		// Try to validate and set the status
		if R.Validate(newPk[:]) {
			tv.Status = StatusValues[Valid]
		} else {
			tv.Status = StatusValues[InvalidPublicKey2]
		}
		tv.Pk2 = hex.EncodeToString(newPk[:])
		tv.Id = n + i
		vectors.Vectors = append(vectors.Vectors, tv)
	}
}

// Private key doesn't correspond to public key
func createNegativePrivateKey(vectors *TestVectors, num int) {
	n := len(vectors.Vectors)
	for i := 0; i < num; i++ {
		var tv TestVector
		var pkBytes1 [csidh.PublicKeySize]byte
		var pkBytes2 [csidh.PublicKeySize]byte
		var prBytes [csidh.PrivateKeySize]byte

		pr1 := R.KeygenPrv()
		pk1 := R.KeygenPub(&pr1)

		// Store private key 1
		pr1.Export(prBytes[:])
		tv.Pr1 = hex.EncodeToString(prBytes[:])

		// Generate public key which doesn't correspond to pr1
		pr2 := R.KeygenPrv()
		pk2 := R.KeygenPub(&pr2)

		pk1.Export(pkBytes1[:])
		pk2.Export(pkBytes2[:])
		if bytes.Equal(pkBytes1[:], pkBytes2[:]) {
			tv.Status = StatusValues[Valid]
		} else {
			tv.Status = StatusValues[InvalidPublicKey1]
		}

		tv.Id = n + i
		tv.Pk1 = hex.EncodeToString(pkBytes2[:])

		vectors.Vectors = append(vectors.Vectors, tv)
	}
}

// TODO: Produce test case for a Pk2 with all zeros and status valid_public_key2.
// comment should be "Zero key should validate correctly"

func main() {
	var vectors TestVectors

	for i := 0; i < 100; i++ {
		vectors.Vectors = append(vectors.Vectors, createValid(i))
	}
	createNegativeSharedSecret(&vectors)
	createNegativePrivateKey(&vectors, 100)
	createNegativePk2(&vectors)

	marshalled, err := json.MarshalIndent(vectors, "", "  ")
	if err != nil {
		panic("Error occured while Marshalling")
	}
	ioutil.WriteFile("testvectors.dat", marshalled, 0644)
}