package e2e

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"math/big"
	"net"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"git.veillax.com/Veillax/blacktip"
	"github.com/emersion/go-msgauth/dkim"
)

// trustedMXCert makes a certificate for 127.0.0.1 and a pool that trusts it, so
// verified TLS to the fake mail server can succeed.
func trustedMXCert(t *testing.T) (*tls.Certificate, *x509.CertPool) {
	key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	tpl := &x509.Certificate{SerialNumber: big.NewInt(7), Subject: pkix.Name{CommonName: "127.0.0.1"}, IPAddresses: []net.IP{net.ParseIP("127.0.0.1")},
		NotBefore: time.Now().Add(-time.Hour), NotAfter: time.Now().Add(time.Hour), KeyUsage: x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		BasicConstraintsValid: true, IsCA: true, ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}}
	der, err := x509.CreateCertificate(rand.Reader, tpl, tpl, &key.PublicKey, key)
	if err != nil {
		t.Fatal(err)
	}
	leaf, _ := x509.ParseCertificate(der)
	pool := x509.NewCertPool()
	pool.AddCert(leaf)
	return &tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key}, pool
}

// asCRLF puts a message back in the form it travels in (the fake server reads it with LF).
func asCRLF(s string) string {
	return strings.ReplaceAll(strings.ReplaceAll(s, "\r\n", "\n"), "\n", "\r\n")
}

type txtTable struct {
	mu sync.Mutex
	m  map[string][]string
}

func (t *txtTable) set(name string, vals ...string) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.m[name] = vals
}

func (t *txtTable) lookup(name string) ([]string, error) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if v, ok := t.m[name]; ok {
		return v, nil
	}
	return nil, &net.DNSError{Err: "no such host", Name: name, IsNotFound: true}
}

func TestDKIMSignsOutgoingMailAndChecksDNS(t *testing.T) {
	mx := newFakeMX(t, false)
	dns := &txtTable{m: map[string][]string{}}
	e := setupMail(t, true, mx, func(c *blacktip.Config) { c.MailLookupTXT = dns.lookup })
	base := "/api/orgs/" + e.org + "/dkim"

	// validation
	for name, in := range map[string]map[string]string{
		"a bad domain":   {"domain": "not a domain"},
		"a wildcard":     {"domain": "*.home.example"},
		"a bad selector": {"domain": "home.example", "selector": "Bad Selector!"},
	} {
		if code, _ := e.admin.do("POST", base, in, nil); code != 400 {
			t.Errorf("%s was accepted", name)
		}
	}
	var key struct {
		ID, Domain, Selector, PublicKey, RecordName, Record, Status string
	}
	e.admin.must("POST", base, map[string]string{"domain": "Home.Example", "selector": "s1"}, &key)
	if key.Domain != "home.example" || key.RecordName != "s1._domainkey.home.example" || !strings.HasPrefix(key.Record, "v=DKIM1; k=rsa; p=") || !strings.HasSuffix(key.Record, key.PublicKey) {
		t.Fatalf("key: %+v", key)
	}
	if code, _ := e.admin.do("POST", base, map[string]string{"domain": "home.example", "selector": "s1"}, nil); code != 400 {
		t.Error("the same selector twice must be refused (rotate by choosing a new one)")
	}
	if code, raw := e.admin.do("GET", base, nil, nil); code != 200 || strings.Contains(raw, "PRIVATE KEY") {
		t.Errorf("listing must work and never reveal the private key: %d %s", code, raw)
	}

	_, listen := e.add(map[string]any{"name": "signed"})
	verify := func(body string) []*dkim.Verification {
		v, err := dkim.VerifyWithOptions(strings.NewReader(asCRLF(body)), &dkim.VerifyOptions{LookupTXT: func(string) ([]string, error) { return []string{key.Record}, nil }})
		if err != nil {
			t.Fatalf("verify: %v", err)
		}
		return v
	}

	// mail from the signed domain is signed, and an independent verifier accepts it
	if err := send(t, listen, nil, "me@home.example", []string{"you@dest.example"}, letter); err != nil {
		t.Fatalf("send: %v", err)
	}
	got := mx.last(t).body
	v := verify(got)
	if len(v) != 1 || v[0].Err != nil || v[0].Domain != "home.example" {
		t.Fatalf("the signature did not verify: %+v", v)
	}
	if !strings.Contains(got, "Hello.") || !strings.Contains(got, ".line that starts with a dot") {
		t.Errorf("signing damaged the body: %q", got)
	}

	// mail from another domain is left alone
	other := strings.ReplaceAll(letter, "me@home.example", "me@other.example")
	if err := send(t, listen, nil, "me@other.example", []string{"you@dest.example"}, other); err != nil {
		t.Fatal(err)
	}
	if b := mx.last(t).body; strings.Contains(b, "DKIM-Signature") {
		t.Errorf("mail for a domain without a key was signed:\n%s", b)
	}

	// a message your mail server already signed for that domain is not signed twice
	pre := "DKIM-Signature: v=1; a=rsa-sha256; d=home.example; s=old; h=from; bh=x; b=y\r\n" + letter
	if err := send(t, listen, nil, "me@home.example", []string{"you@dest.example"}, pre); err != nil {
		t.Fatal(err)
	}
	if n := strings.Count(strings.ToLower(mx.last(t).body), "dkim-signature:"); n != 1 {
		t.Errorf("an already-signed message now has %d signatures", n)
	}

	// subdomains are covered by the parent domain's key
	sub := strings.ReplaceAll(letter, "me@home.example", "me@news.home.example")
	if err := send(t, listen, nil, "me@news.home.example", []string{"you@dest.example"}, sub); err != nil {
		t.Fatal(err)
	}
	if v := verify(mx.last(t).body); len(v) != 1 || v[0].Err != nil {
		t.Errorf("a subdomain wasn't signed by the parent's key: %+v", v)
	}

	// the DNS check: missing until you publish the record
	status := func() string {
		var list []struct{ ID, Status string }
		e.admin.must("GET", base, nil, &list)
		return list[0].Status
	}
	waitStatus2 := func(want string) {
		deadline := time.Now().Add(12 * time.Second)
		for status() != want {
			if time.Now().After(deadline) {
				t.Fatalf("DKIM status stuck at %q, want %q", status(), want)
			}
			time.Sleep(250 * time.Millisecond)
		}
	}
	waitStatus2("missing")
	dns.set("s1._domainkey.home.example", "v=DKIM1; k=rsa; p=AAAA") // somebody else's key
	e.admin.must("PUT", "/api/orgs/"+e.org+"/mail-egress/"+firstMailEgress(t, e), egressBody(e, "nudge1.example.invalid"), nil)
	waitStatus2("mismatch")
	dns.set("s1._domainkey.home.example", "v=DKIM1; k=rsa; p="+key.PublicKey[:40], "p="+key.PublicKey[40:]) // published, split over strings like long records are
	dns.set("s1._domainkey.home.example", "v=DKIM1; k=rsa; "+"p="+key.PublicKey)
	e.admin.must("PUT", "/api/orgs/"+e.org+"/mail-egress/"+firstMailEgress(t, e), egressBody(e, "nudge2.example.invalid"), nil)
	waitStatus2("published")

	// removing the key stops signing
	e.admin.must("DELETE", base+"/"+key.ID, nil, nil)
	deadline := time.Now().Add(8 * time.Second)
	for {
		if err := send(t, listen, nil, "me@home.example", []string{"you@dest.example"}, letter); err != nil {
			t.Fatal(err)
		}
		if !strings.Contains(mx.last(t).body, "DKIM-Signature") {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("mail is still signed after the key was deleted")
		}
		time.Sleep(300 * time.Millisecond)
	}
}

func firstMailEgress(t *testing.T, e *mailEnv) string {
	var st struct{ MailEgress []struct{ ID string } }
	e.admin.must("GET", "/api/orgs/"+e.org+"/state", nil, &st)
	return st.MailEgress[0].ID
}

func egressBody(e *mailEnv, helo string) map[string]any {
	var st struct {
		MailEgress []struct {
			Listen string
		}
	}
	e.admin.must("GET", "/api/orgs/"+e.org+"/state", nil, &st)
	return map[string]any{"name": "signed", "siteId": e.site, "edgeId": e.edge, "sourceIp": "127.0.0.2", "heloName": helo, "listen": st.MailEgress[0].Listen, "enabled": true}
}

func TestMTASTSAndVerifiedTLS(t *testing.T) {
	cert, pool := trustedMXCert(t)
	good := newFakeMX(t, true) // offers STARTTLS with a certificate the edge trusts
	good.cert = cert
	plain := newFakeMX(t, false)
	untrusted := newFakeMX(t, true) // self-signed, not trusted

	var mu sync.Mutex
	policies := map[string]*blacktip.STSPolicy{}
	e := setupMail(t, true, good, func(c *blacktip.Config) {
		c.MailRootCAs = pool
		c.MailSTS = func(domain string) (*blacktip.STSPolicy, error) {
			mu.Lock()
			defer mu.Unlock()
			return policies[domain], nil
		}
	})
	hostOf := func(m *fakeMX) blacktip.MailHost {
		h, p, _ := net.SplitHostPort(m.addr)
		n, _ := strconv.Atoi(p)
		return blacktip.MailHost{Host: h, Port: n}
	}
	e.mu.Lock()
	e.hosts["tls.example"] = []blacktip.MailHost{hostOf(good)}
	e.hosts["plain.example"] = []blacktip.MailHost{hostOf(plain)}
	e.hosts["untrusted.example"] = []blacktip.MailHost{hostOf(untrusted)}
	e.mu.Unlock()
	mu.Lock()
	policies["tls.example"] = &blacktip.STSPolicy{Mode: "enforce", MX: []string{"127.0.0.1"}}
	policies["plain.example"] = &blacktip.STSPolicy{Mode: "enforce", MX: []string{"127.0.0.1"}}
	policies["untrusted.example"] = &blacktip.STSPolicy{Mode: "enforce", MX: []string{"127.0.0.1"}}
	policies["wrongmx.example"] = &blacktip.STSPolicy{Mode: "enforce", MX: []string{"mail.elsewhere.example", "*.mx.elsewhere.example"}}
	policies["testing.example"] = &blacktip.STSPolicy{Mode: "testing", MX: []string{"mail.elsewhere.example"}}
	mu.Unlock()
	e.mu.Lock()
	e.hosts["wrongmx.example"] = []blacktip.MailHost{hostOf(good)}
	e.hosts["testing.example"] = []blacktip.MailHost{hostOf(plain)}
	e.mu.Unlock()
	_, listen := e.add(map[string]any{"name": "sts"})

	mail := func(domain string) error {
		return send(t, listen, nil, "me@home.example", []string{"you@" + domain}, letter)
	}
	// enforce + a listed server + verified TLS: delivered, over TLS
	if err := mail("tls.example"); err != nil {
		t.Fatalf("enforce with a good server: %v", err)
	}
	if got := good.last(t); !got.tls {
		t.Error("the message was not sent over TLS under an enforcing policy")
	}
	// enforce + a server the policy doesn't list: refused before any connection
	if err := mail("wrongmx.example"); smtpCode(err) != 451 || !strings.Contains(err.Error(), "MTA-STS") {
		t.Errorf("a server outside the policy: %v", err)
	}
	// enforce + no TLS offered: refused, and nothing sent in the clear
	if err := mail("plain.example"); smtpCode(err) != 451 || !strings.Contains(err.Error(), "TLS") || plain.count() != 0 {
		t.Errorf("enforce without TLS: err=%v delivered=%d", err, plain.count())
	}
	// enforce + a certificate nobody trusts: refused
	if err := mail("untrusted.example"); smtpCode(err) != 451 || untrusted.count() != 0 {
		t.Errorf("enforce with an untrusted certificate: err=%v delivered=%d", err, untrusted.count())
	}
	// testing mode only observes: delivered even in the clear
	if err := mail("testing.example"); err != nil {
		t.Errorf("testing mode must not block mail: %v", err)
	}
	// and with no policy at all, mail flows as before
	if err := send(t, listen, nil, "me@home.example", []string{"you@dest.example"}, letter); err != nil {
		t.Errorf("no policy: %v", err)
	}

	// the egress's own "require verified TLS" now succeeds against a trusted certificate and still refuses the rest
	_, strict := e.add(map[string]any{"name": "strict", "requireTls": true})
	if err := send(t, strict, nil, "me@home.example", []string{"you@tls.example"}, letter); err != nil {
		t.Errorf("require TLS against a trusted server: %v", err)
	}
	if err := send(t, strict, nil, "me@home.example", []string{"you@untrusted.example"}, letter); smtpCode(err) != 451 {
		t.Errorf("require TLS against an untrusted server: %v", err)
	}
}