//go:build ignore

// Prints the Go reference values of src/lib/dkc/ibe.ts (plan of phase 2,
// section 4) as the JSON of src/lib/dkc/testing/ibe-vectors.json. Everything
// comes from the libraries that tlock decrypts with for Quicknet
// (bls-unchained-g1-rfc9380): drand/kyber encrypt/ibe on
// kyber-bls12381.NewBLS12381Suite(), over kilic/bls12-381.
//
//   - gt: e(G1, G2) and e(2·G1, G2), serialized by kyber-bls12381 (the order
//     of kilic: c1 before c0 at every level of the tower), with H2 truncated
//     to 16 and 32 bytes.
//   - h3 and h4: H3 and H4 on fixed inputs, among them inputs whose first
//     candidates for r are rejected.
//   - round_identities: the identity of a round, scheme.DigestBeacon.
//   - fixtures: for the tlock stanza of every official .dkc, the pairing of
//     the release signature with U, sigma, r and the file key. The file key
//     is the one tlock.TimeUnlock unwraps, and age.Decrypt opens
//     OUTER_TIME_AGE with it: the header MAC and the STREAM verify.
//   - kyber: messages of 0, 1, 16 and 32 bytes encrypted for round 1000 by
//     ibe.EncryptCCAonG2 itself, with its random sigma, and decrypted back by
//     ibe.DecryptCCAonG2.
//   - decrypt: the verdict of ibe.DecryptCCAonG2, after decoding the points
//     as the scheme does, on edited copies of the time_only stanza.
//
// H2, H3 and H4 are unexported in kyber: this file restates them with
// kyber's exported tags, and checks them on every fixture against what
// tlock.TimeUnlock unwraps and against U = r·G2. A mismatch panics.
//
// The kyber vectors use a random sigma, so the output is not
// byte-reproducible: the file is generated once and frozen, like the .dkc
// fixtures of the reference.
//
// 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 directory of the official fixtures:
//
//	go run ibe-go-vectors.go path/to/testdata/fixtures > ibe-vectors.json
package main

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

	"filippo.io/age"
	"g.activething.com/go/DateKeys/capsule"
	"g.activething.com/go/DateKeys/profile"
	"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 signature of round 1001, as in the mutation corpus.
const sig1001 = "b33bf3667cbd5a82de3a24b4e0e9fe5513cc1a0e840368c6e31f5fcfa79bea03f73896b25883abf2853d10337fb8fa41"

var (
	suite = bls.NewBLS12381Suite()
	// The field and the scalar orders of BLS12-381.
	p, _     = new(big.Int).SetString("1a0111ea397fe69a4b1ba7b6434bacd764774b84f38512bf6730d2a0f6b0f6241eabfffeb153ffffb9feffffffffaaab", 16)
	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 marshal(m interface{ MarshalBinary() ([]byte, error) }) string {
	return hex.EncodeToString(must(m.MarshalBinary()))
}

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
}

// H2, H3 and H4 of kyber encrypt/ibe (gtToHash, h3, h4), restated.
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]
}

// h3 returns r as 32 big-endian bytes and the iteration that accepted it.
func h3(sigma, msg []byte) ([]byte, int) {
	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 {
			return d[:], int(i)
		}
	}
	panic("h3: rejection sampling failed")
}

// rG2 is r·G2 through kyber, with r decoded as kyber's h3 decodes it.
func rG2(r []byte) kyber.Point {
	s := suite.G2().Scalar()
	if err := s.UnmarshalBinary(r); err != nil {
		panic(err)
	}
	return suite.G2().Point().Mul(s, nil)
}

type gtVector struct {
	Name string `json:"name"`
	G1   string `json:"g1"`
	G2   string `json:"g2"`
	GT   string `json:"gt"`
	H2   string `json:"h2_16"`
	H232 string `json:"h2_32"`
}

type h3Vector struct {
	Name       string `json:"name"`
	Sigma      string `json:"sigma"`
	Msg        string `json:"msg"`
	R          string `json:"r"`
	Iterations int    `json:"iterations"`
}

type h4Vector struct {
	Sigma string `json:"sigma"`
	H416  string `json:"h4_16"`
	H432  string `json:"h4_32"`
}

type roundIdentity struct {
	Round uint64 `json:"round"`
	ID    string `json:"id"`
}

type fixtureVector struct {
	Name      string `json:"name"`
	Round     uint64 `json:"round"`
	Signature string `json:"signature"`
	Body      string `json:"body"`
	GT        string `json:"gt"`
	Sigma     string `json:"sigma"`
	R         string `json:"r"`
	FileKey   string `json:"file_key"`
}

type ciphertextVector struct {
	Name      string `json:"name"`
	Round     uint64 `json:"round"`
	Signature string `json:"signature"`
	U         string `json:"u"`
	V         string `json:"v"`
	W         string `json:"w"`
	Go        string `json:"go"`
	Msg       string `json:"msg,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)
	}
	out := struct {
		Description     string             `json:"description"`
		Generator       string             `json:"generator"`
		Libraries       string             `json:"libraries"`
		Scheme          string             `json:"scheme"`
		PublicKey       string             `json:"public_key"`
		GT              []gtVector         `json:"gt"`
		H3              []h3Vector         `json:"h3"`
		H4              []h4Vector         `json:"h4"`
		RoundIdentities []roundIdentity    `json:"round_identities"`
		Fixtures        []fixtureVector    `json:"fixtures"`
		Kyber           []ciphertextVector `json:"kyber"`
		Decrypt         []ciphertextVector `json:"decrypt"`
	}{
		Description: "Go reference values of the tlock IBE-CCA on G2 (spec §63 step 11) for src/lib/dkc/ibe.ts; " +
			"see scripts/ibe-go-vectors.go for how each block is obtained.",
		Generator: "scripts/ibe-go-vectors.go",
		Libraries: libraries(),
		Scheme:    scheme.Name,
		PublicKey: hex.EncodeToString(quicknet.PublicKey),
	}

	// GT and H2.
	g1, g2 := suite.G1().Point().Base(), suite.G2().Point().Base()
	two := suite.G1().Point().Mul(suite.G1().Scalar().SetInt64(2), g1)
	square := suite.GT().Point().Add(suite.Pair(g1, g2), suite.Pair(g1, g2))
	if !square.Equal(suite.Pair(two, g2)) {
		panic("e(2·G1, G2) is not e(G1, G2) squared")
	}
	for _, c := range []struct {
		name string
		a, b kyber.Point
	}{{"e(G1, G2)", g1, g2}, {"e(2·G1, G2), the square of e(G1, G2)", two, g2}} {
		gt := must(suite.Pair(c.a, c.b).MarshalBinary())
		out.GT = append(out.GT, gtVector{c.name, marshal(c.a), marshal(c.b), hex.EncodeToString(gt),
			hex.EncodeToString(h2(gt, 16)), hex.EncodeToString(h2(gt, 32))})
	}

	// H3: fixed inputs, then the first inputs of a deterministic sequence
	// whose r is accepted at the second and at the third iteration.
	for _, c := range []struct{ name, sigma, msg string }{
		{"16 zero bytes each", strings.Repeat("00", 16), strings.Repeat("00", 16)},
		{"32 bytes each", strings.Repeat("ab", 32), strings.Repeat("cd", 32)},
		{"empty", "", ""},
	} {
		r, it := h3(unhex(c.sigma), unhex(c.msg))
		out.H3 = append(out.H3, h3Vector{c.name, c.sigma, c.msg, hex.EncodeToString(r), it})
	}
	for want := 2; want <= 3; want++ {
		for i := uint32(0); ; i++ {
			seed := sha256.Sum256(binary.BigEndian.AppendUint32([]byte("DateKeys H3 vector "), i))
			sigma, msg := seed[:16], seed[16:]
			if r, it := h3(sigma, msg); it == want {
				out.H3 = append(out.H3, h3Vector{fmt.Sprintf("accepted at iteration %d (sequence item %d)", want, i),
					hex.EncodeToString(sigma), hex.EncodeToString(msg), hex.EncodeToString(r), it})
				break
			}
		}
	}

	// H4.
	for _, sigma := range []string{strings.Repeat("00", 16), strings.Repeat("5a", 32)} {
		out.H4 = append(out.H4, h4Vector{sigma, hex.EncodeToString(h4(unhex(sigma), 16)), hex.EncodeToString(h4(unhex(sigma), 32))})
	}

	// Round identities.
	for _, round := range []uint64{1, 1000, 1001, 83903165811, 1<<53 - 1} {
		out.RoundIdentities = append(out.RoundIdentities, roundIdentity{round, hex.EncodeToString(scheme.DigestBeacon(&common.Beacon{Round: round}))})
	}

	// The fixtures.
	dir := os.Args[1]
	names := must(filepath.Glob(filepath.Join(dir, "*.dkc")))
	sort.Strings(names)
	var timeOnly fixtureVector
	for _, name := range names {
		v := fixture(scheme, key, name)
		out.Fixtures = append(out.Fixtures, v)
		if v.Name == "time_only" {
			timeOnly = v
		}
	}
	if timeOnly.Name == "" {
		panic("no time_only fixture")
	}

	// Encryptions by kyber, for round 1000.
	id1000 := scheme.DigestBeacon(&common.Beacon{Round: 1000})
	sig1000 := unhex(timeOnly.Signature)
	for _, n := range []int{0, 1, 16, 32} {
		msg := sha256.Sum256([]byte("DateKeys IBE vector message"))
		ct := must(ibe.EncryptCCAonG2(suite, key, id1000, msg[:n]))
		v := verdict(scheme, fmt.Sprintf("a %d-byte message", n), 1000, sig1000, unhex(marshal(ct.U)), ct.V, ct.W)
		if v.Go != "ok" || v.Msg != hex.EncodeToString(msg[:n]) {
			panic("kyber does not decrypt its own ciphertext")
		}
		out.Kyber = append(out.Kyber, v)
	}

	// Edited copies of the time_only stanza.
	body := unhex(timeOnly.Body)
	u, vv, w := body[:96], body[96:112], body[112:]
	flip := func(b []byte, i int, mask byte) []byte {
		c := bytes.Clone(b)
		c[i] ^= mask
		return c
	}
	// c0 is the second coordinate of the compressed encoding of G2.
	c0 := new(big.Int).SetBytes(u[48:])
	uc0p := concat(u[:48], new(big.Int).Add(c0, p).FillBytes(make([]byte, 48)))
	infinityG2 := concat([]byte{0xc0}, make([]byte, 95))
	infinityG1 := concat([]byte{0xc0}, make([]byte, 47))
	for _, c := range []struct {
		name         string
		sig, u, v, w []byte
	}{
		{"the time_only stanza", sig1000, u, vv, w},
		{"U with p added to c0", sig1000, uc0p, vv, w},
		{"U is the point at infinity", sig1000, infinityG2, vv, w},
		{"U negated", sig1000, flip(u, 0, 0x20), vv, w},
		{"V with its first bit flipped", sig1000, u, flip(vv, 0, 0x80), w},
		{"W with its last bit flipped", sig1000, u, vv, flip(w, 15, 0x01)},
		{"the signature of round 1001", unhex(sig1001), u, vv, w},
		{"the signature negated", flip(sig1000, 0, 0x20), u, vv, w},
		{"the signature is the point at infinity", infinityG1, u, vv, w},
		{"W one byte shorter than V", sig1000, u, vv, w[:15]},
		{"V and W of 33 bytes", sig1000, u, concat(vv, vv, []byte{0}), concat(w, w, []byte{0})},
	} {
		out.Decrypt = append(out.Decrypt, verdict(scheme, c.name, 1000, c.sig, c.u, c.v, c.w))
	}

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

// fixture opens the OUTER_TIME_AGE of a .dkc with the release of its sidecar,
// through tlock.TimeUnlock inside an age identity, and restates the IBE.
func fixture(scheme *crypto.Scheme, key kyber.Point, path string) fixtureVector {
	file := must(os.ReadFile(path))
	var side struct {
		Release struct {
			Round     uint64 `json:"round"`
			Signature string `json:"signature"`
		} `json:"release"`
	}
	if err := json.Unmarshal(must(os.ReadFile(strings.TrimSuffix(path, ".dkc")+".json")), &side); err != nil {
		panic(err)
	}
	sig := unhex(side.Release.Signature)
	pre := must(capsule.ParsePrelude(file))
	start := capsule.PreludeSize + int(pre.PublicHeaderLen)
	sealed := file[start : start+int(pre.SealedControlLen)]

	var body, fileKey []byte
	id := unwrap(func(stanzas []*age.Stanza) ([]byte, error) {
		if len(stanzas) != 1 || stanzas[0].Type != "tlock" || len(stanzas[0].Args) != 2 || stanzas[0].Args[0] != strconv.FormatUint(side.Release.Round, 10) {
			panic("not one tlock stanza for the round of the release")
		}
		body = stanzas[0].Body
		ct := must(tlock.BytesToCiphertext(*scheme, body))
		fileKey = must(tlock.TimeUnlock(*scheme, key, common.Beacon{Round: side.Release.Round, Signature: sig}, ct))
		return bytes.Clone(fileKey), nil
	})
	r := must(age.Decrypt(bytes.NewReader(sealed), id))
	if _, err := io.Copy(io.Discard, r); err != nil {
		panic(fmt.Sprintf("%s: the age payload does not open: %v", path, err))
	}

	// The IBE restated, checked against tlock.
	u, v, w := body[:96], body[96:112], body[112:]
	sp, up := scheme.SigGroup.Point(), scheme.KeyGroup.Point()
	if sp.UnmarshalBinary(sig) != nil || up.UnmarshalBinary(u) != nil {
		panic("points do not decode")
	}
	gt := must(suite.Pair(sp, up).MarshalBinary())
	sigma := xor(v, h2(gt, len(w)))
	msg := xor(w, h4(sigma, len(w)))
	if !bytes.Equal(msg, fileKey) {
		panic(path + ": the restated H2 and H4 disagree with tlock")
	}
	rb, _ := h3(sigma, msg)
	if !rG2(rb).Equal(up) {
		panic(path + ": the restated H3 disagrees with U")
	}
	return fixtureVector{strings.TrimSuffix(filepath.Base(path), ".dkc"), side.Release.Round, side.Release.Signature,
		hex.EncodeToString(body), hex.EncodeToString(gt), hex.EncodeToString(sigma), hex.EncodeToString(rb), hex.EncodeToString(fileKey)}
}

// verdict decodes the points as the scheme does and runs ibe.DecryptCCAonG2,
// the decryption of tlock.TimeUnlock for Quicknet after its beacon check.
func verdict(scheme *crypto.Scheme, name string, round uint64, sig, u, v, w []byte) ciphertextVector {
	out := ciphertextVector{Name: name, Round: round, Signature: hex.EncodeToString(sig), U: hex.EncodeToString(u),
		V: hex.EncodeToString(v), W: hex.EncodeToString(w), Go: "reject"}
	sp, up := scheme.SigGroup.Point(), scheme.KeyGroup.Point()
	if sp.UnmarshalBinary(sig) != nil || up.UnmarshalBinary(u) != nil {
		return out
	}
	msg, err := ibe.DecryptCCAonG2(suite, sp, &ibe.Ciphertext{U: up, V: v, W: w})
	if err != nil {
		return out
	}
	out.Go, out.Msg = "ok", hex.EncodeToString(msg)
	return out
}

type unwrap func([]*age.Stanza) ([]byte, error)

func (f unwrap) Unwrap(stanzas []*age.Stanza) ([]byte, error) { return f(stanzas) }

// 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, ", ")
}