package capsule_test

import (
	"bytes"
	"context"
	"crypto/elliptic"
	"crypto/sha256"
	"encoding/hex"
	"g.activething.com/go/DateKeys/internal/cms/cmstest"
	"io"
	"strings"
	"testing"
	"time"

	"g.activething.com/go/DateKeys/authorkey"
	"g.activething.com/go/DateKeys/capsule"
	"g.activething.com/go/DateKeys/internal/testkit"
)

// openSigned opens dkc with the author keys that the person saved.
func openSigned(t *testing.T, dkc []byte, saved map[string]string) *capsule.Opened {
	t.Helper()
	o := defaultOpen(1000)
	o.Sink, o.AuthorKeys = &testkit.MemorySink{}, saved
	opened, err := capsule.Open(context.Background(), nil, bytes.NewReader(dkc), o)
	if err != nil {
		t.Fatal(err)
	}
	return opened
}

// Spec v0.11 §29.7, §29.8, §62.1 rule 19: EncryptFiles signs with the key of
// opts.AuthorKey and Open gives F4, or F3 with the key saved.
func TestEncryptFilesSigned(t *testing.T) {
	key, err := authorkey.Generate()
	if err != nil {
		t.Fatal(err)
	}
	pub, _ := authorkey.PublicString(key.Public())
	opts := files3(t)
	opts.AuthorKey = key
	var dkc bytes.Buffer
	if _, err := capsule.EncryptFiles(&dkc, []capsule.Source{source("nota.txt", "Hola.\n")}, opts); err != nil {
		t.Fatal(err)
	}

	o := openSigned(t, dkc.Bytes(), nil)
	if o.Verdicts.Signature != capsule.VerdictSignedOther || o.Verdicts.AuthorKey != [32]byte(key.Public()) || o.Verdicts.Seal != capsule.VerdictNoSeal || o.AreaLen != capsule.AreaLen {
		t.Errorf("verdicts %+v, area %d", o.Verdicts, o.AreaLen)
	}
	o = openSigned(t, dkc.Bytes(), map[string]string{pub: "Ana"})
	if o.Verdicts.Signature != capsule.VerdictSignedSaved || o.Verdicts.AuthorLabel != "Ana" {
		t.Errorf("saved key: verdicts %+v", o.Verdicts)
	}
	other, _ := authorkey.Generate()
	otherPub, _ := authorkey.PublicString(other.Public())
	if o = openSigned(t, dkc.Bytes(), map[string]string{otherPub: "Luis"}); o.Verdicts.Signature != capsule.VerdictSignedOther {
		t.Errorf("another saved key: verdicts %+v", o.Verdicts)
	}

	// LargeArea only allows widening: a signature that fits keeps the area of 32 KiB.
	opts.LargeArea = true
	dkc.Reset()
	if _, err := capsule.EncryptFiles(&dkc, []capsule.Source{source("nota.txt", "Hola.\n")}, opts); err != nil {
		t.Fatal(err)
	}
	if o = openSigned(t, dkc.Bytes(), nil); o.AreaLen != capsule.AreaLen || o.Verdicts.Signature != capsule.VerdictSignedOther {
		t.Errorf("large area: verdicts %+v, area %d", o.Verdicts, o.AreaLen)
	}
}

// badKey is an AuthorKey that signs wrongly or has a public key of the wrong
// length.
type badKey struct {
	pub, sig []byte
}

func (k badKey) Public() []byte     { return k.pub }
func (k badKey) Sign([]byte) []byte { return k.sig }

// Spec v0.11 §62.1 rule 19: a signature that does not verify, or a key of
// another length, fails before anything is written.
func TestEncryptFilesSignatureChecked(t *testing.T) {
	key, _ := authorkey.Generate()
	for _, tc := range []struct {
		name string
		key  capsule.AuthorKey
		want string
	}{
		{"wrong signature", badKey{key.Public(), make([]byte, 64)}, "self-check"},
		{"short key", badKey{key.Public()[:31], make([]byte, 64)}, "not 32"},
	} {
		opts := files3(t)
		opts.AuthorKey = tc.key
		var dkc bytes.Buffer
		_, err := capsule.EncryptFiles(&dkc, []capsule.Source{source("nota.txt", "x")}, opts)
		if err == nil || !strings.Contains(err.Error(), tc.want) {
			t.Errorf("%s: %v", tc.name, err)
		}
		if dkc.Len() != 0 {
			t.Errorf("%s: %d bytes written", tc.name, dkc.Len())
		}
	}
}

// Spec v0.11 §29.7, §29.8: a signature of the fixture format3_signed holds
// in the context of its capsule and in no other: a bit of the signature, of
// a commitment or of the message changes the verdict to F2; the same message
// signed by another key is another key's F4; and a security area without it
// is F0.
func TestSignedFixtureVerdicts(t *testing.T) {
	f := loadFixture(t, "format3_signed")
	body := f.plaintext
	frame, err := capsule.ParseBodyFrame(body[:capsule.BodyFrameSize], uint64(len(body)))
	if err != nil {
		t.Fatal(err)
	}
	security := body[capsule.BodyFrameSize : capsule.BodyFrameSize+frame.SecurityLen]
	cb, _ := hex.DecodeString(f.ControlCBOR)
	c, err := capsule.DecodeControl(cb, f.format())
	if err != nil {
		t.Fatal(err)
	}
	cc, err := capsule.ControlCommit(c, f.format())
	if err != nil {
		t.Fatal(err)
	}
	hb := body[capsule.BodyFrameSize+frame.AreaLen : capsule.BodyFrameSize+frame.AreaLen+frame.HeadLen]
	ctx := func() *capsule.SecurityContext {
		return &capsule.SecurityContext{ControlCommit: cc, HeadDigest: capsule.HeadDigest(hb)}
	}
	if v := capsule.EvaluateSecurityIn(security, ctx()); v.Signature != capsule.VerdictSignedOther || v.Seal != capsule.VerdictNoSeal {
		t.Fatalf("the signature of the fixture: %+v", v)
	}

	// Another capsule: another control commitment, or another head.
	other := ctx()
	other.ControlCommit[0] ^= 1
	if v := capsule.EvaluateSecurityIn(security, other); v.Signature != capsule.VerdictSignatureInvalid {
		t.Errorf("another control: %+v", v)
	}
	other = ctx()
	other.HeadDigest[31] ^= 1
	if v := capsule.EvaluateSecurityIn(security, other); v.Signature != capsule.VerdictSignatureInvalid {
		t.Errorf("another head: %+v", v)
	}

	// A bit of the signature, or of the key.
	_, value, err := capsule.SecurityKey2(security)
	if err != nil {
		t.Fatal(err)
	}
	pub, _ := authorkey.ParsePublic(f.Signature.AuthorKey)
	for name, mutate := range map[string]func(sig, key []byte){
		"signature": func(sig, key []byte) { sig[63] ^= 1 },
		"key":       func(sig, key []byte) { key[0] ^= 1 },
	} {
		sig, key := bytes.Clone(value), bytes.Clone(pub)
		mutate(sig, key)
		x, err := capsule.EncodeAuthorSignature(capsule.AlgEd25519, key, sig)
		if err != nil {
			t.Fatal(err)
		}
		s, err := capsule.EncodeSecurityWith(x, nil)
		if err != nil {
			t.Fatal(err)
		}
		if v := capsule.EvaluateSecurityIn(s, ctx()); v.Signature != capsule.VerdictSignatureInvalid {
			t.Errorf("a bit of the %s: %+v", name, v)
		}
	}

	// The same message signed by another key: F4 with that key.
	msg := capsule.AuthorMessage(cc, capsule.HeadDigest(hb), capsule.SignersDigest(capsule.AlgEd25519, nil))
	k, _ := authorkey.Generate()
	x, _ := capsule.EncodeAuthorSignature(capsule.AlgEd25519, k.Public(), k.Sign(msg))
	s, _ := capsule.EncodeSecurityWith(x, nil)
	if v := capsule.EvaluateSecurityIn(s, ctx()); v.Signature != capsule.VerdictSignedOther || v.AuthorKey != [32]byte(k.Public()) {
		t.Errorf("another key: %+v", v)
	}

	// Removed: F0.
	if v := capsule.EvaluateSecurityIn(capsule.EncodeSecurity(), ctx()); v.Signature != capsule.VerdictNoSignature {
		t.Errorf("removed: %+v", v)
	}
}

// cmsSigner is a CMSSigner that signs as a signing application would: its
// certificates sign AUTHOR_MESSAGE, and the authority tsa seals each
// signature at when.
type cmsSigner struct {
	signers []cmstest.Signer
	tsa     cmstest.Signer
	when    time.Time
	seen    []byte // the message it was asked to sign
	calls   int
	junk    int // bytes of an unsigned attribute that decides nothing
}

func (c *cmsSigner) Signers() (out [][32]byte) {
	for _, s := range c.signers {
		out = append(out, sha256.Sum256(s.Cert.Raw))
	}
	return out
}

func (c *cmsSigner) Sign(message []byte) ([]byte, error) {
	c.seen = message
	c.calls++
	opts := cmstest.Options{Junk: c.junk}
	if c.tsa.Key != nil {
		opts.Token = func(sig []byte) []byte {
			return cmstest.Token(sig, c.when, cmstest.TokenOptions{Accuracy: time.Second}, c.tsa)
		}
	}
	return cmstest.Signature(message, opts, c.signers...), nil
}

// sealer is a Sealer that asks the authority tsa.
type sealer struct {
	tsa  cmstest.Signer
	when time.Time
}

func (s sealer) Seal(subject [32]byte) ([]byte, error) {
	return cmstest.Token(subject[:], s.when, cmstest.TokenOptions{}, s.tsa), nil
}

// Spec v0.11 §29.10, §29.11, §62.1 rules 19 and 21: EncryptFiles gives
// AUTHOR_MESSAGE to the CMSSigner, checks that what it returns is complete,
// and Open gives F6; a Sealer seals SEAL_SUBJECT and Open gives S4.
func TestEncryptFilesCMSAndSeal(t *testing.T) {
	from, to := time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC), time.Date(2040, 1, 1, 0, 0, 0, 0, time.UTC)
	ana := cmstest.NewECDSA("Ana López", elliptic.P256(), from, to)
	luis := cmstest.NewRSA("Luis Gómez", 2048, from, to)
	tsa := cmstest.NewECDSA("TSA de prueba", elliptic.P256(), from, to)
	opts := files3(t)
	when := opts.Now()

	signer := &cmsSigner{signers: []cmstest.Signer{ana, luis}, tsa: tsa, when: when}
	opts.CMSSigner = signer
	var dkc bytes.Buffer
	if _, err := capsule.EncryptFiles(&dkc, []capsule.Source{source("nota.txt", "Hola.\n")}, opts); err != nil {
		t.Fatal(err)
	}
	o := openSigned(t, dkc.Bytes(), nil)
	if o.Verdicts.Signature != capsule.VerdictSignedComplete || o.Verdicts.Seal != capsule.VerdictNoSeal || len(o.Verdicts.Detail.Signers) != 2 ||
		len(signer.seen) != capsule.AuthorMessageSize || capsule.AuthorCode(signer.seen) == "" {
		t.Errorf("verdicts %+v, message %q", o.Verdicts, signer.seen)
	}
	if !o.Verdicts.Detail.Signers[0].Before {
		t.Error("the seal does not precede the round time")
	}

	// A signature that lacks a required signer is not written.
	third := cmstest.NewECDSA("Falta", elliptic.P256(), from, to)
	missing := &cmsSigner{signers: []cmstest.Signer{ana}, tsa: tsa, when: when}
	opts.CMSSigner = &requiring{missing, third}
	dkc.Reset()
	if _, err := capsule.EncryptFiles(&dkc, []capsule.Source{source("nota.txt", "x")}, opts); err == nil || !strings.Contains(err.Error(), "F5") || dkc.Len() != 0 {
		t.Errorf("a missing signer: %v, %d bytes written", err, dkc.Len())
	}
	// Without seals the signature is incomplete too.
	opts.CMSSigner = &cmsSigner{signers: []cmstest.Signer{ana}}
	if _, err := capsule.EncryptFiles(&dkc, []capsule.Source{source("nota.txt", "x")}, opts); err == nil || !strings.Contains(err.Error(), "without seal") {
		t.Errorf("no seal: %v", err)
	}

	// A seal over the signature of an author key.
	key, _ := authorkey.Generate()
	opts.CMSSigner, opts.AuthorKey, opts.Sealer = nil, key, sealer{tsa, when}
	dkc.Reset()
	if _, err := capsule.EncryptFiles(&dkc, []capsule.Source{source("nota.txt", "Hola.\n")}, opts); err != nil {
		t.Fatal(err)
	}
	o = openSigned(t, dkc.Bytes(), nil)
	if o.Verdicts.Signature != capsule.VerdictSignedOther || o.Verdicts.Seal != capsule.VerdictSealed || o.Verdicts.Detail.SealHolder != "TSA de prueba" {
		t.Errorf("verdicts %+v", o.Verdicts)
	}
	// And a seal over a capsule without a signature.
	opts.AuthorKey = nil
	dkc.Reset()
	if _, err := capsule.EncryptFiles(&dkc, []capsule.Source{source("nota.txt", "Hola.\n")}, opts); err != nil {
		t.Fatal(err)
	}
	if o = openSigned(t, dkc.Bytes(), nil); o.Verdicts.Signature != capsule.VerdictNoSignature || o.Verdicts.Seal != capsule.VerdictSealed {
		t.Errorf("verdicts %+v", o.Verdicts)
	}

	// The exclusions.
	opts.AuthorKey, opts.CMSSigner = key, signer
	if _, err := capsule.EncryptFiles(io.Discard, []capsule.Source{source("a", "x")}, opts); err == nil {
		t.Error("AuthorKey and CMSSigner")
	}
	opts.AuthorKey = nil
	if _, err := capsule.EncryptFiles(io.Discard, []capsule.Source{source("a", "x")}, opts); err == nil {
		t.Error("CMSSigner and Sealer")
	}
}

// requiring asks for one signer more than signs.
type requiring struct {
	*cmsSigner
	extra cmstest.Signer
}

func (r *requiring) Signers() [][32]byte {
	return append(r.cmsSigner.Signers(), sha256.Sum256(r.extra.Cert.Raw))
}

// Review: what is signed does not depend on the area, so the area is chosen
// after the signature, and widening it never makes anybody sign twice. A
// signature that fits keeps the common area, LargeArea or not.
func TestAreaChosenAfterSigning(t *testing.T) {
	from, to := time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC), time.Date(2040, 1, 1, 0, 0, 0, 0, time.UTC)
	ana := cmstest.NewECDSA("Ana López", elliptic.P256(), from, to)
	tsa := cmstest.NewECDSA("TSA de prueba", elliptic.P256(), from, to)
	opts := files3(t)
	when := opts.Now()

	// 40 KB of signature: too much for 32 KiB, and the person signs once.
	big := &cmsSigner{signers: []cmstest.Signer{ana}, tsa: tsa, when: when, junk: 40 << 10}
	opts.CMSSigner = big
	if _, err := capsule.EncryptFiles(io.Discard, []capsule.Source{source("nota.txt", "x")}, opts); err == nil || !strings.Contains(err.Error(), "LargeArea") || big.calls != 1 {
		t.Errorf("without LargeArea: %v after %d calls", err, big.calls)
	}
	opts.LargeArea = true
	big.calls = 0
	var dkc bytes.Buffer
	res, err := capsule.EncryptFiles(&dkc, []capsule.Source{source("nota.txt", "Hola.\n")}, opts)
	if err != nil || big.calls != 1 {
		t.Fatalf("with LargeArea: %v after %d calls", err, big.calls)
	}
	o := openSigned(t, dkc.Bytes(), nil)
	if o.AreaLen != capsule.LargeAreaLen || o.Verdicts.Signature != capsule.VerdictSignedComplete || o.PayloadLength != res.Length || o.PaddedLength != res.PaddedLength {
		t.Errorf("area %d, verdicts %+v, L %d P %d", o.AreaLen, o.Verdicts, o.PayloadLength, o.PaddedLength)
	}
	// An unsigned capsule with LargeArea keeps the common area: P does not
	// tell that someone asked.
	plain := files3(t)
	plain.LargeArea = true
	dkc.Reset()
	if _, err := capsule.EncryptFiles(&dkc, []capsule.Source{source("nota.txt", "x")}, plain); err != nil {
		t.Fatal(err)
	}
	if o := openSigned(t, dkc.Bytes(), nil); o.AreaLen != capsule.AreaLen {
		t.Errorf("an unsigned capsule with LargeArea: area %d", o.AreaLen)
	}
}

// Review: a typed nil is an error, never a capsule without the signature or
// the seal that was asked for; the exclusions are checked before a file is
// read, and format 2 refuses the options it cannot honour.
func TestWriterOptionsChecked(t *testing.T) {
	for _, set := range []func(*capsule.EncryptOptions){
		func(o *capsule.EncryptOptions) { o.AuthorKey = (*authorkey.Key)(nil) },
		func(o *capsule.EncryptOptions) { o.CMSSigner = (*cmsSigner)(nil) },
		func(o *capsule.EncryptOptions) { o.Sealer = (*sealer)(nil) },
	} {
		opts := files3(t)
		set(&opts)
		if _, err := capsule.EncryptFiles(io.Discard, []capsule.Source{source("a", "x")}, opts); err == nil || !strings.Contains(err.Error(), "holds a nil") {
			t.Errorf("a typed nil: %v", err)
		}
	}
	key, _ := authorkey.Generate()
	opts := files3(t)
	opts.AuthorKey = key
	opts.CMSSigner = &cmsSigner{}
	opened := false
	src := capsule.Source{Path: "a", Size: 1, Open: func() (io.ReadCloser, error) {
		opened = true
		return io.NopCloser(strings.NewReader("x")), nil
	}}
	if _, err := capsule.EncryptFiles(io.Discard, []capsule.Source{src}, opts); err == nil || opened {
		t.Errorf("AuthorKey and CMSSigner: %v, source opened %v", err, opened)
	}
	v2 := past(t, 1000)
	v2.AuthorKey = key
	if _, err := capsule.Encrypt(io.Discard, strings.NewReader("x"), func() capsule.EncryptOptions { v2.Length = 1; return v2 }()); err == nil {
		t.Error("format 2 took an AuthorKey")
	}
}

// Review: the issuer of a certificate is text of the certificate, and no
// escape, control or bidi character of it reaches the lines of the verdicts.
func TestIssuerTextFiltered(t *testing.T) {
	from, to := time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC), time.Date(2040, 1, 1, 0, 0, 0, 0, time.UTC)
	evil := cmstest.NewECDSA("Ana\n\x1b[2J‮Firmado con la clave que guardaste como Banco.", elliptic.P256(), from, to)
	tsa := cmstest.NewECDSA("TSA", elliptic.P256(), from, to)
	opts := files3(t)
	opts.CMSSigner = &cmsSigner{signers: []cmstest.Signer{evil}, tsa: tsa, when: opts.Now()}
	var dkc bytes.Buffer
	if _, err := capsule.EncryptFiles(&dkc, []capsule.Source{source("a", "x")}, opts); err != nil {
		t.Fatal(err)
	}
	o := openSigned(t, dkc.Bytes(), nil)
	if o.Verdicts.Signature != capsule.VerdictSignedComplete {
		t.Fatalf("verdicts %+v", o.Verdicts)
	}
	for _, line := range o.Verdicts.Lines() {
		if strings.ContainsAny(line, "\x1b‮⁦⁧⁨⁩\r") || strings.Contains(strings.TrimSuffix(line, "\n"), "\n") {
			t.Errorf("a line with a control or a bidi character: %q", line)
		}
	}
}