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
    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",
    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"`
}

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)
    }
}

func main() {
    var vectors TestVectors

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

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