//go:build ignore

// Prints src/lib/dkc/testing/tlock-vectors.json: the Go reference values for
// the tlock encryption of src/lib/dkc/ibe.ts and src/lib/dkc/tlock.ts (plan
// of phase 2, section 8, points 4 and 5).
//
//   - encrypt: EncryptCCAonG2 of drand/kyber with the suite of tlock for
//     Quicknet, restated with a fixed sigma, for messages of 1, 16 and 32
//     bytes to rounds 1000 and 1001. kyber draws sigma from crypto/rand, so
//     the restatement is checked: ibe.DecryptCCAonG2 opens every ciphertext
//     with the published signature of its round, and tlock.BytesToCiphertext
//     with tlock.TimeUnlock opens the 16-byte ones, as a tlock stanza body.
//   - interop: the ciphertexts that scripts/tlock-ts-samples.mjs made with
//     the TypeScript library, each opened by the reference. An IBE body goes
//     through tlock.BytesToCiphertext and tlock.TimeUnlock; an age file
//     through age.Decrypt with agewrap.NewTimeIdentity, the identity of
//     capsule.Open at step 11. Each sample gets the verdict "ok" with what Go
//     recovered, or "reject".
//
// H2, H3 and H4 are unexported in kyber; they are restated with its tags, as
// in scripts/ibe-go-vectors.go, and the decryptions above check them.
//
// Run it from a scratch module that requires the reference implementation
// (replace g.activething.com/go/DateKeys => ../datekeys-go and
// GOFLAGS=-mod=mod), passing the output of tlock-ts-samples.mjs:
//
//	go run tlock-go-vectors.go ts-samples.json > tlock-vectors.json
package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"math/big"
	"os"
	"runtime/debug"
	"sort"
	"strings"

	"filippo.io/age"
	"g.activething.com/go/DateKeys/agewrap"
	"g.activething.com/go/DateKeys/profile"
	"g.activething.com/go/DateKeys/provider"
	"github.com/drand/drand/v2/common"
	"github.com/drand/drand/v2/crypto"
	"github.com/drand/kyber"
	bls "github.com/drand/kyber-bls12381"
	"github.com/drand/kyber/encrypt/ibe"
	"github.com/drand/tlock"
)

// The published Quicknet signatures of rounds 1000 and 1001, as in the
// fixtures and the mutation corpus; provider.Verify checks them below.
var published = map[uint64]string{
	1000: "b44679b9a59af2ec876b1a6b1ad52ea9b1615fc3982b19576350f93447cb1125e342b73a8dd2bacbe47e4b6b63ed5e39",
	1001: "b33bf3667cbd5a82de3a24b4e0e9fe5513cc1a0e840368c6e31f5fcfa79bea03f73896b25883abf2853d10337fb8fa41",
}

var (
	suite    = bls.NewBLS12381Suite()
	order, _ = new(big.Int).SetString("73eda753299d7d483339d80809a1d80553bda402fffe5bfeffffffff00000001", 16)
)

func must[T any](v T, err error) T {
	if err != nil {
		panic(err)
	}
	return v
}

func unhex(s string) []byte { return must(hex.DecodeString(s)) }

func concat(parts ...[]byte) []byte { return bytes.Join(parts, nil) }

func xor(a, b []byte) []byte {
	out := make([]byte, len(a))
	for i := range a {
		out[i] = a[i] ^ b[i]
	}
	return out
}

func h2(gt []byte, n int) []byte {
	sum := sha256.Sum256(concat(ibe.H2Tag(), gt))
	return sum[:n]
}

func h4(sigma []byte, n int) []byte {
	sum := sha256.Sum256(concat(ibe.H4Tag(), sigma))
	return sum[:n]
}

func h3(sigma, msg []byte) kyber.Scalar {
	base := sha256.Sum256(concat(ibe.H3Tag(), sigma, msg))
	for i := uint16(1); i < 65535; i++ {
		d := sha256.Sum256(concat(binary.LittleEndian.AppendUint16(nil, i), base[:]))
		d[0] >>= 1
		if new(big.Int).SetBytes(d[:]).Cmp(order) < 0 {
			r := suite.G1().Scalar()
			if err := r.UnmarshalBinary(d[:]); err != nil {
				panic(err)
			}
			return r
		}
	}
	panic("h3: rejection sampling failed")
}

// encrypt is EncryptCCAonG2 of kyber with the given sigma.
func encrypt(key kyber.Point, id, msg, sigma []byte) *ibe.Ciphertext {
	qid := suite.G1().Point().(kyber.HashablePoint).Hash(id)
	gid := suite.Pair(qid, key)
	r := h3(sigma, msg)
	gt := must(suite.GT().Point().Mul(r, gid).MarshalBinary())
	return &ibe.Ciphertext{U: suite.G2().Point().Mul(r, nil), V: xor(sigma, h2(gt, len(msg))), W: xor(msg, h4(sigma, len(msg)))}
}

type encryptVector struct {
	Name      string `json:"name"`
	Round     uint64 `json:"round"`
	ID        string `json:"id"`
	Msg       string `json:"msg"`
	Sigma     string `json:"sigma"`
	U         string `json:"u"`
	V         string `json:"v"`
	W         string `json:"w"`
	Signature string `json:"signature"`
}

type sample struct {
	Name      string `json:"name"`
	Kind      string `json:"kind"`
	Round     uint64 `json:"round"`
	FileKey   string `json:"file_key,omitempty"`
	Body      string `json:"body,omitempty"`
	Plaintext string `json:"plaintext,omitempty"`
	File      string `json:"file,omitempty"`
	Go        string `json:"go"`
	GoResult  string `json:"go_result,omitempty"`
}

func main() {
	scheme := must(crypto.SchemeFromName(crypto.SigsOnG1ID))
	quicknet := profile.Quicknet()
	key := scheme.KeyGroup.Point()
	if err := key.UnmarshalBinary(quicknet.PublicKey); err != nil {
		panic(err)
	}
	release := func(round uint64) provider.Release {
		r := provider.Release{Round: round, Signature: unhex(published[round])}
		if err := provider.Verify(quicknet, provider.Condition{Round: round}, r); err != nil {
			panic(err)
		}
		return r
	}

	var vectors []encryptVector
	for i, c := range []struct {
		round uint64
		n     int
	}{{1000, 16}, {1001, 16}, {1000, 1}, {1000, 32}} {
		seed := sha256.Sum256(binary.BigEndian.AppendUint32([]byte("DateKeys tlock vector "), uint32(i)))
		msgSeed := sha256.Sum256(seed[:])
		msg, sigma := msgSeed[:c.n], seed[:c.n]
		id := scheme.DigestBeacon(&common.Beacon{Round: c.round})
		ct := encrypt(key, id, msg, sigma)
		rel := release(c.round)
		sig := scheme.SigGroup.Point()
		if err := sig.UnmarshalBinary(rel.Signature); err != nil {
			panic(err)
		}
		if got, err := ibe.DecryptCCAonG2(suite, sig, ct); err != nil || !bytes.Equal(got, msg) {
			panic(fmt.Sprintf("kyber does not decrypt the restated encryption %d: %v", i, err))
		}
		u := must(ct.U.MarshalBinary())
		if c.n == 16 {
			body := concat(u, ct.V, ct.W)
			got := must(tlock.TimeUnlock(*scheme, key, common.Beacon{Round: rel.Round, Signature: rel.Signature}, must(tlock.BytesToCiphertext(*scheme, body))))
			if !bytes.Equal(got, msg) {
				panic("tlock does not decrypt the restated encryption")
			}
		}
		vectors = append(vectors, encryptVector{
			Name: fmt.Sprintf("a %d-byte message for round %d", c.n, c.round), Round: c.round, ID: hex.EncodeToString(id),
			Msg: hex.EncodeToString(msg), Sigma: hex.EncodeToString(sigma), U: hex.EncodeToString(u),
			V: hex.EncodeToString(ct.V), W: hex.EncodeToString(ct.W), Signature: published[c.round],
		})
	}

	var in struct {
		Generator string   `json:"generator"`
		Samples   []sample `json:"samples"`
	}
	if err := json.Unmarshal(must(os.ReadFile(os.Args[1])), &in); err != nil {
		panic(err)
	}
	for i := range in.Samples {
		s := &in.Samples[i]
		rel := release(s.Round)
		s.Go = "reject"
		switch s.Kind {
		case "ibe":
			ct, err := tlock.BytesToCiphertext(*scheme, unhex(s.Body))
			if err != nil {
				fmt.Fprintf(os.Stderr, "%s: %v\n", s.Name, err)
				continue
			}
			fk, err := tlock.TimeUnlock(*scheme, key, common.Beacon{Round: rel.Round, Signature: rel.Signature}, ct)
			if err != nil {
				fmt.Fprintf(os.Stderr, "%s: %v\n", s.Name, err)
				continue
			}
			s.Go, s.GoResult = "ok", hex.EncodeToString(fk)
		case "age":
			id := must(agewrap.NewTimeIdentity(quicknet, s.Round, rel))
			r, err := age.Decrypt(bytes.NewReader(unhex(s.File)), id)
			if err != nil {
				fmt.Fprintf(os.Stderr, "%s: %v\n", s.Name, err)
				continue
			}
			pt, err := io.ReadAll(r)
			if err != nil {
				fmt.Fprintf(os.Stderr, "%s: %v\n", s.Name, err)
				continue
			}
			s.Go, s.GoResult = "ok", hex.EncodeToString(pt)
		default:
			panic("unknown sample kind " + s.Kind)
		}
	}

	out := map[string]any{
		"description": "Go reference values for the tlock encryption of src/lib/dkc/ibe.ts and src/lib/dkc/tlock.ts; " +
			"see scripts/tlock-go-vectors.go for how each block is obtained.",
		"generator":  "scripts/tlock-go-vectors.go",
		"libraries":  libraries(),
		"scheme":     scheme.Name,
		"public_key": hex.EncodeToString(quicknet.PublicKey),
		"encrypt":    vectors,
		"interop":    map[string]any{"generator": in.Generator, "samples": in.Samples},
	}
	enc := json.NewEncoder(os.Stdout)
	enc.SetIndent("", " ")
	enc.SetEscapeHTML(false)
	if err := enc.Encode(out); err != nil {
		panic(err)
	}
}

// libraries names the versions of the libraries this program ran with.
func libraries() string {
	info, ok := debug.ReadBuildInfo()
	if !ok {
		panic("no build info")
	}
	var out []string
	for _, d := range info.Deps {
		switch d.Path {
		case "filippo.io/age", "github.com/drand/tlock", "github.com/drand/kyber", "github.com/drand/kyber-bls12381", "github.com/drand/drand/v2":
			out = append(out, d.Path+" "+d.Version)
		}
	}
	sort.Strings(out)
	return strings.Join(out, ", ")
}