package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"go/parser"
	"go/token"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"

	"filippo.io/age"
	bls "github.com/drand/kyber-bls12381"
)

const (
	fixturesDir = "../../testdata/fixtures"
	releasesDir = "../../testdata/releases"
	vectorsDir  = "../../testdata/vectors"
)

type fixtureRecord struct {
	Format  int `json:"format"`
	Release struct {
		Round     uint64 `json:"round"`
		Signature string `json:"signature"`
	} `json:"release"`
	PayloadLength   uint64 `json:"payload_length"`
	PlaintextFile   string `json:"plaintext_file"`
	PlaintextSHA256 string `json:"plaintext_sha256"`
	AccessKeyFile   string `json:"access_key_file"`
	Files           []struct {
		Path   string `json:"path"`
		Size   uint64 `json:"size"`
		SHA256 string `json:"sha256"`
	} `json:"files"`
}

func readJSON(t *testing.T, path string, v any) {
	t.Helper()
	b, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	if err := json.Unmarshal(b, v); err != nil {
		t.Fatalf("%s: %v", path, err)
	}
}

func unhex(t *testing.T, s string) []byte {
	t.Helper()
	b, err := hex.DecodeString(s)
	if err != nil {
		t.Fatal(err)
	}
	return b
}

// A tiny deterministic CBOR encoder, enough for a release object.
func cborHead(major byte, n uint64) []byte {
	switch {
	case n < 24:
		return []byte{major<<5 | byte(n)}
	case n < 1<<8:
		return []byte{major<<5 | 24, byte(n)}
	case n < 1<<16:
		return []byte{major<<5 | 25, byte(n >> 8), byte(n)}
	case n < 1<<32:
		return []byte{major<<5 | 26, byte(n >> 24), byte(n >> 16), byte(n >> 8), byte(n)}
	}
	return []byte{major<<5 | 27, byte(n >> 56), byte(n >> 48), byte(n >> 40), byte(n >> 32),
		byte(n >> 24), byte(n >> 16), byte(n >> 8), byte(n)}
}

func encodeRelease(chainHash []byte, round uint64, sig []byte) []byte {
	var b []byte
	b = append(b, cborHead(5, 5)...)
	b = append(b, cborHead(0, 0)...)
	b = append(b, cborHead(3, 16)...)
	b = append(b, "datekeys-release"...)
	b = append(b, cborHead(0, 1)...)
	b = append(b, cborHead(0, 1)...)
	b = append(b, cborHead(0, 2)...)
	b = append(b, cborHead(2, uint64(len(chainHash)))...)
	b = append(b, chainHash...)
	b = append(b, cborHead(0, 3)...)
	b = append(b, cborHead(0, round)...)
	b = append(b, cborHead(0, 4)...)
	b = append(b, cborHead(2, uint64(len(sig)))...)
	b = append(b, sig...)
	return b
}

// committedRelease reads testdata/releases/<round>.cbor and checks that it is
// the release object built here from the fixture record.
func committedRelease(t *testing.T, round uint64, sigHex string) []byte {
	t.Helper()
	want := encodeRelease(unhex(t, quicknetChainHash), round, unhex(t, sigHex))
	path := filepath.Join(releasesDir, strconv.FormatUint(round, 10)+".cbor")
	got, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("the committed release is missing: %v", err)
	}
	if !bytes.Equal(got, want) {
		t.Fatalf("%s:\n got %x\nwant %x", path, got, want)
	}
	return got
}

func loadFixture(t *testing.T, name string) (fixtureRecord, []byte, []byte, []byte) {
	t.Helper()
	var rec fixtureRecord
	readJSON(t, filepath.Join(fixturesDir, name+".json"), &rec)
	dkc, err := os.ReadFile(filepath.Join(fixturesDir, name+".dkc"))
	if err != nil {
		t.Fatal(err)
	}
	var dkk []byte
	if rec.AccessKeyFile != "" {
		if dkk, err = os.ReadFile(filepath.Join(fixturesDir, rec.AccessKeyFile)); err != nil {
			t.Fatal(err)
		}
	}
	return rec, dkc, committedRelease(t, rec.Release.Round, rec.Release.Signature), dkk
}

func TestRecoverFixtures(t *testing.T) {
	for _, name := range []string{
		"format3_single",
		"format3_time_and_key_portable",
		"format3_tree",
		"format3_signed",
		"format2_time_only", // two STREAM chunks
		"format2_time_and_key_portable",
		"format2_time_and_key_sixteen", // round 2000; no .dkk, so only its release is checked below
		"time_only",
		"time_and_key_portable",
	} {
		t.Run(name, func(t *testing.T) {
			rec, dkc, relObj, dkk := loadFixture(t, name)
			if name == "format2_time_and_key_sixteen" {
				_, err := recoverCapsule(dkc, relObj, nil, io.Discard)
				if err == nil || !strings.Contains(err.Error(), "needs its .dkk") {
					t.Fatalf("got %v, want the .dkk to be required", err)
				}
				return
			}
			res, err := recoverCapsule(dkc, relObj, dkk, io.Discard)
			if err != nil {
				t.Fatal(err)
			}
			want, err := os.ReadFile(filepath.Join(fixturesDir, rec.PlaintextFile))
			if err != nil {
				t.Fatal(err)
			}
			if !bytes.Equal(res.content, want) {
				t.Fatalf("content differs from %s", rec.PlaintextFile)
			}
			if res.format != rec.Format || uint64(len(res.content)) != rec.PayloadLength {
				t.Fatalf("format %d, L %d; want %d, %d", res.format, len(res.content), rec.Format, rec.PayloadLength)
			}
			if res.format != 3 {
				return
			}
			if len(res.body.files) != len(rec.Files) {
				t.Fatalf("%d files, want %d", len(res.body.files), len(rec.Files))
			}
			dir := t.TempDir()
			if err := writeFiles(dir, res.body, io.Discard); err != nil {
				t.Fatal(err)
			}
			for i, f := range rec.Files {
				if res.body.files[i].path != f.Path {
					t.Fatalf("file %d is %q, want %q", i, res.body.files[i].path, f.Path)
				}
				b, err := os.ReadFile(filepath.Join(dir, filepath.FromSlash(f.Path)))
				if err != nil {
					t.Fatal(err)
				}
				sum := sha256.Sum256(b)
				if hex.EncodeToString(sum[:]) != f.SHA256 || uint64(len(b)) != f.Size {
					t.Fatalf("%s: written file differs from the record", f.Path)
				}
			}
		})
	}
}

func TestRun(t *testing.T) {
	dir := t.TempDir()
	out, bodyFile := filepath.Join(dir, "out"), filepath.Join(dir, "body")
	args := []string{
		"-dkc", filepath.Join(fixturesDir, "format3_time_and_key_portable.dkc"),
		"-release", filepath.Join(releasesDir, "1000.cbor"),
		"-dkk", filepath.Join(fixturesDir, "format3_time_and_key_portable.dkk"),
		"-out", out, "-body", bodyFile,
	}
	var log bytes.Buffer
	if err := run(args, &log); err != nil {
		t.Fatalf("%v\n%s", err, log.String())
	}
	got, _ := os.ReadFile(bodyFile)
	want, _ := os.ReadFile(filepath.Join(fixturesDir, "format3_time_and_key_portable.plaintext"))
	if !bytes.Equal(got, want) {
		t.Fatal("-body differs from the plaintext fixture")
	}
	if _, err := os.Stat(filepath.Join(out, "secreto.txt")); err != nil {
		t.Fatal(err)
	}
}

func TestBadReleases(t *testing.T) {
	var single, other fixtureRecord
	readJSON(t, filepath.Join(fixturesDir, "format3_single.json"), &single)
	readJSON(t, filepath.Join(fixturesDir, "empty_payload.json"), &other) // round 1001
	dkc, err := os.ReadFile(filepath.Join(fixturesDir, "format3_single.dkc"))
	if err != nil {
		t.Fatal(err)
	}
	chain := unhex(t, quicknetChainHash)
	wrongChain := bytes.Clone(chain)
	wrongChain[0] ^= 1
	sig, otherSig := unhex(t, single.Release.Signature), unhex(t, other.Release.Signature)

	for _, tc := range []struct {
		name, want string
		relObj     []byte
	}{
		{"wrong chain_hash", "not Quicknet's", encodeRelease(wrongChain, 1000, sig)},
		{"release of another round", "needs round 1000", encodeRelease(chain, other.Release.Round, otherSig)},
		{"signature of another round", "does not verify", encodeRelease(chain, 1000, otherSig)},
		{"wrong type tag", "type tag", bytes.Replace(encodeRelease(chain, 1000, sig), []byte("release"), []byte("relaxed"), 1)},
	} {
		t.Run(tc.name, func(t *testing.T) {
			_, err := recoverCapsule(dkc, tc.relObj, nil, io.Discard)
			if err == nil || !strings.Contains(err.Error(), tc.want) {
				t.Fatalf("got %v, want an error with %q", err, tc.want)
			}
		})
	}
}

// TestTlockVectors checks steps 10 and 11 value by value against
// testdata/vectors/tlock_steps.json and the GT serialization against
// tlock_ibe.json.
func TestTlockVectors(t *testing.T) {
	var ibe struct {
		Vectors []struct{ G1, G2, GT, H2 string }
	}
	readJSON(t, filepath.Join(vectorsDir, "tlock_ibe.json"), &ibe)
	for _, v := range ibe.Vectors {
		g1, g2 := bls.NullKyberG1(), bls.NullKyberG2()
		if err := decodePoint(g1, unhex(t, v.G1), 48, "g1"); err != nil {
			t.Fatal(err)
		}
		if err := decodePoint(g2, unhex(t, v.G2), 96, "g2"); err != nil {
			t.Fatal(err)
		}
		if !g1.Equal(bls.NullKyberG1().Base()) || !g2.Equal(bls.NullKyberG2().Base()) {
			t.Fatal("vector points are not the generators")
		}
		gt := suite.Pair(g1, g2)
		raw, _ := gt.MarshalBinary()
		if hex.EncodeToString(raw) != v.GT {
			t.Fatal("GT serialization differs")
		}
		if h, _ := h2(gt); hex.EncodeToString(h) != v.H2 {
			t.Fatalf("H2 = %x, want %s", h, v.H2)
		}
	}

	var steps struct {
		ChainHash string `json:"chain_hash"`
		PublicKey string `json:"public_key"`
		DST       string `json:"dst"`
		Vectors   []struct {
			Round                           uint64
			Signature, Message              string
			HashToG1                        string `json:"hash_to_g1"`
			Body, Pairing, H2, Sigma, H4, R string
			FileKey                         string `json:"file_key"`
		}
	}
	readJSON(t, filepath.Join(vectorsDir, "tlock_steps.json"), &steps)
	if steps.ChainHash != quicknetChainHash || steps.PublicKey != quicknetPublicKey || steps.DST != quicknetDST {
		t.Fatal("pinned parameters differ from the vectors")
	}
	for _, v := range steps.Vectors {
		if hex.EncodeToString(roundMessage(v.Round)) != v.Message {
			t.Fatalf("round %d: M differs", v.Round)
		}
		hm, _ := hashToG1(roundMessage(v.Round)).MarshalBinary()
		if hex.EncodeToString(hm) != v.HashToG1 {
			t.Fatalf("round %d: H(M) differs", v.Round)
		}
		sig, err := verifyRelease(release{round: v.Round, signature: unhex(t, v.Signature)})
		if err != nil {
			t.Fatal(err)
		}
		body := unhex(t, v.Body)
		u := bls.NullKyberG2()
		if err := decodePoint(u, body[:96], 96, "U"); err != nil {
			t.Fatal(err)
		}
		gt := suite.Pair(sig, u)
		raw, _ := gt.MarshalBinary()
		mask, _ := h2(gt)
		sigma := xor(body[96:112], mask)
		r, err := h3(sigma, xor(body[112:], h4(sigma)))
		if err != nil {
			t.Fatal(err)
		}
		fk, err := unwrapTlock(body, sig)
		if err != nil {
			t.Fatal(err)
		}
		for _, c := range []struct{ name, got, want string }{
			{"pairing", hex.EncodeToString(raw), v.Pairing},
			{"H2", hex.EncodeToString(mask), v.H2},
			{"sigma", hex.EncodeToString(sigma), v.Sigma},
			{"H4", hex.EncodeToString(h4(sigma)), v.H4},
			{"r", hex.EncodeToString(r.FillBytes(make([]byte, 32))), v.R},
			{"file key", hex.EncodeToString(fk), v.FileKey},
		} {
			if c.got != c.want {
				t.Fatalf("round %d: %s = %s, want %s", v.Round, c.name, c.got, c.want)
			}
		}
	}
}

func TestPinnedRoundTime(t *testing.T) {
	// format3_single.json: round 1000 unlocks at 2023-08-23T15:59:24Z.
	if got := roundTime(1000).Format("2006-01-02T15:04:05Z"); got != "2023-08-23T15:59:24Z" {
		t.Fatal(got)
	}
}

// TestAgeSecretKey checks the Bech32 encoder against age's own.
func TestAgeSecretKey(t *testing.T) {
	var rec struct {
		AccessMaterial string `json:"access_material"`
	}
	readJSON(t, filepath.Join(fixturesDir, "format3_time_and_key_portable.dkk.json"), &rec)
	key := ageSecretKey(unhex(t, rec.AccessMaterial))
	got, err := age.ParseX25519Identity(key)
	if err != nil {
		t.Fatal(err)
	}
	if got.String() != key {
		t.Fatalf("age re-encodes %s as %s", key, got.String())
	}
}

func TestStreamEdges(t *testing.T) {
	if _, err := streamDecrypt(make([]byte, 32), nil); err == nil {
		t.Fatal("an empty STREAM must fail")
	}
}

// TestTampered flips one byte in each part of a capsule: the tlock U, the
// sealed age payload and PAYLOAD_AGE must each make the recovery fail.
func TestTampered(t *testing.T) {
	var rec fixtureRecord
	readJSON(t, filepath.Join(fixturesDir, "format3_single.json"), &rec)
	dkc, err := os.ReadFile(filepath.Join(fixturesDir, "format3_single.dkc"))
	if err != nil {
		t.Fatal(err)
	}
	relObj := encodeRelease(unhex(t, quicknetChainHash), rec.Release.Round, unhex(t, rec.Release.Signature))
	sealedEnd := 16 + 0x79 + 0x1ca // PUBLIC_HEADER_LEN and SEALED_CONTROL_LEN of its prelude
	stanza := bytes.Index(dkc, []byte("-> tlock"))
	body := stanza + bytes.IndexByte(dkc[stanza:], '\n') + 1
	for _, tc := range []struct {
		name string
		at   int
	}{
		{"tlock U", body + 10},
		{"sealed control STREAM", sealedEnd - 1},
		{"PAYLOAD_AGE STREAM", len(dkc) - 1},
	} {
		t.Run(tc.name, func(t *testing.T) {
			bad := bytes.Clone(dkc)
			bad[tc.at] ^= 0x01
			if _, err := recoverCapsule(bad, relObj, nil, io.Discard); err == nil {
				t.Fatal("a tampered capsule opened")
			}
		})
	}
}

// TestImports keeps the program independent of DateKeys and of the drand
// tools: only the standard library, age, kyber and x/crypto.
func TestImports(t *testing.T) {
	f, err := parser.ParseFile(token.NewFileSet(), "main.go", nil, parser.ImportsOnly)
	if err != nil {
		t.Fatal(err)
	}
	var errs []error
	for _, imp := range f.Imports {
		p, _ := strconv.Unquote(imp.Path.Value)
		for _, banned := range []string{"g.activething.com/go/DateKeys", "github.com/drand/tlock", "github.com/drand/drand"} {
			if strings.HasPrefix(p, banned) {
				errs = append(errs, errors.New("forbidden import "+p))
			}
		}
	}
	if err := errors.Join(errs...); err != nil {
		t.Fatal(err)
	}
}