// generates ML-DSA test vectors using Go's crypto/internal/fips140/mldsa
// run from within the Go source tree because mldsa is not public yet
// see https://github.com/golang/go/issues/77626
//
// this is normally ran via gen.vsh
//
// GOROOT=. ./bin/go run ./src/crypto/internal/fips140/mldsa/gen
package main

import (
    "crypto/internal/fips140/mldsa"
    "crypto/internal/fips140/sha3"
    "crypto/sha256"
    "encoding/hex"
    "encoding/json"
    "fmt"
    "os"
)

type Vector struct {
    Kind      string `json:"kind"`
    Seed      string `json:"seed"`
    Msg       string `json:"msg"`
    PkSha256  string `json:"pk_sha256"`
    SigSha256 string `json:"sig_sha256"`
    Context   string `json:"context,omitempty"`
}

type variant struct {
    name          string
    newPrivateKey func([]byte) (*mldsa.PrivateKey, error)
    newPublicKey  func([]byte) (*mldsa.PublicKey, error)
}

func main() {
    variants := []variant{
        {"ml_dsa_44", mldsa.NewPrivateKey44, mldsa.NewPublicKey44},
        {"ml_dsa_65", mldsa.NewPrivateKey65, mldsa.NewPublicKey65},
        {"ml_dsa_87", mldsa.NewPrivateKey87, mldsa.NewPublicKey87},
    }

    s := sha3.NewShake128()
    seed := make([]byte, 32)
    var vectors []Vector

    for _, v := range variants {
        for i := 0; i < 3; i++ {
            s.Read(seed)
            priv, err := v.newPrivateKey(seed)
            if err != nil {
                panic(err)
            }
            pk := priv.PublicKey().Bytes()

            msg := make([]byte, 32+i*17)
            s.Read(msg)

            sig, err := mldsa.SignDeterministic(priv, msg, "")
            if err != nil {
                panic(err)
            }

            pub, err := v.newPublicKey(pk)
            if err != nil {
                panic(err)
            }
            if err := mldsa.Verify(pub, msg, sig, ""); err != nil {
                panic(fmt.Sprintf("verify failed: %v", err))
            }

            pkHash := sha256.Sum256(pk)
            sigHash := sha256.Sum256(sig)
            vectors = append(vectors, Vector{
                Kind:      v.name,
                Seed:      hex.EncodeToString(seed),
                Msg:       hex.EncodeToString(msg),
                PkSha256:  hex.EncodeToString(pkHash[:]),
                SigSha256: hex.EncodeToString(sigHash[:]),
            })
        }

        s.Read(seed)
        priv, err := v.newPrivateKey(seed)
        if err != nil {
            panic(err)
        }
        pk := priv.PublicKey().Bytes()
        msg := make([]byte, 40)
        s.Read(msg)

        sig, err := mldsa.SignDeterministic(priv, msg, "test-context")
        if err != nil {
            panic(err)
        }
        pub, err := v.newPublicKey(pk)
        if err != nil {
            panic(err)
        }
        if err := mldsa.Verify(pub, msg, sig, "test-context"); err != nil {
            panic(fmt.Sprintf("verify with context failed: %v", err))
        }

        pkHash := sha256.Sum256(pk)
        sigHash := sha256.Sum256(sig)
        vectors = append(vectors, Vector{
            Kind:      v.name,
            Seed:      hex.EncodeToString(seed),
            Msg:       hex.EncodeToString(msg),
            PkSha256:  hex.EncodeToString(pkHash[:]),
            SigSha256: hex.EncodeToString(sigHash[:]),
            Context:   "test-context",
        })
    }

    enc := json.NewEncoder(os.Stdout)
    enc.SetIndent("", "  ")
    if err := enc.Encode(vectors); err != nil {
        panic(err)
    }
}