package encdns

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"io"
	"math/big"
	"net"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/miekg/dns"
)

func testCert(t *testing.T, names ...string) (tls.Certificate, *x509.Certificate) {
	key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	tmpl := &x509.Certificate{SerialNumber: big.NewInt(0xabc123), Subject: pkix.Name{CommonName: names[0]}, DNSNames: names,
		NotBefore: time.Now().Add(-time.Hour), NotAfter: time.Now().Add(24 * time.Hour), ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}}
	der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
	if err != nil {
		t.Fatal(err)
	}
	leaf, _ := x509.ParseCertificate(der)
	return tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key, Leaf: leaf}, leaf
}

// answer answers every question with 192.0.2.53.
var answer = dns.HandlerFunc(func(w dns.ResponseWriter, r *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(r)
	m.Answer = append(m.Answer, &dns.A{Hdr: dns.RR_Header{Name: r.Question[0].Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 60}, A: net.ParseIP("192.0.2.53")})
	w.WriteMsg(m)
})

func listeners(t *testing.T, cert tls.Certificate) (dot, doh Endpoint) {
	cfg := &tls.Config{Certificates: []tls.Certificate{cert}}
	ln, err := tls.Listen("tcp", "127.0.0.1:0", cfg)
	if err != nil {
		t.Fatal(err)
	}
	srv := &dns.Server{Listener: ln, Net: "tcp-tls", Handler: answer}
	go srv.ActivateAndServe()
	t.Cleanup(func() { srv.Shutdown() })
	hl, err := tls.Listen("tcp", "127.0.0.1:0", cfg)
	if err != nil {
		t.Fatal(err)
	}
	hs := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/dns-query" {
			http.NotFound(w, r)
			return
		}
		b, _ := io.ReadAll(r.Body)
		q := new(dns.Msg)
		if q.Unpack(b) != nil {
			http.Error(w, "bad", 400)
			return
		}
		m := new(dns.Msg)
		m.SetReply(q)
		m.Rcode = dns.RcodeSuccess
		out, _ := m.Pack()
		w.Header().Set("Content-Type", "application/dns-message")
		w.Write(out)
	})}
	go hs.Serve(hl)
	t.Cleanup(func() { hs.Close() })
	port := func(l net.Listener) int { return l.Addr().(*net.TCPAddr).Port }
	return Endpoint{"127.0.0.1", port(ln)}, Endpoint{"127.0.0.1", port(hl)}
}

func TestHealth(t *testing.T) {
	ctx := context.Background()
	cert, leaf := testCert(t, "dns.example.com")
	dot, doh := listeners(t, cert)
	roots := x509.NewCertPool()
	roots.AddCert(leaf)
	st := Status{Enabled: true, DoT: []Endpoint{dot}, DoH: []Endpoint{doh}}

	h := Check(ctx, st, "dns.example.com", "abc123", roots)
	if !h.OK() || len(h.Probes) != 2 {
		t.Fatalf("working setup: %+v", h.Probes)
	}
	for _, p := range h.Probes {
		if !p.Trusted || p.RCode != "NOERROR" || p.Warning != "" || p.CertSerial != "abc123" {
			t.Errorf("%s: %+v", p.Kind, p)
		}
	}
	// named hasn't reread a renewed certificate.
	if h := Check(ctx, st, "dns.example.com", "def456", roots); h.OK() || !strings.Contains(h.Probes[0].Problem, "older certificate") {
		t.Errorf("stale: %+v", h.Probes[0])
	}
	// The wrong name.
	if h := Check(ctx, st, "other.example.com", "", roots); h.OK() || !strings.Contains(h.Probes[0].Problem, "not other.example.com") {
		t.Errorf("wrong name: %+v", h.Probes[0])
	}
	// A CA this server doesn't trust: works, with a warning.
	if h := Check(ctx, st, "dns.example.com", "abc123", x509.NewCertPool()); !h.OK() || !strings.Contains(h.Probes[1].Warning, "staging") {
		t.Errorf("untrusted: %+v", h.Probes[1])
	}
	// Nothing listening.
	l, _ := net.Listen("tcp", "127.0.0.1:0")
	gone := l.Addr().(*net.TCPAddr).Port
	l.Close()
	if h := Check(ctx, Status{Enabled: true, DoT: []Endpoint{{"127.0.0.1", gone}}}, "dns.example.com", "", roots); h.OK() || !strings.Contains(h.Probes[0].Problem, "nothing listens") {
		t.Errorf("nothing listening: %+v", h.Probes[0])
	}
}

func TestRequireNames(t *testing.T) {
	h := &Health{Probes: []Probe{
		{Kind: "dot", OK: true, CertNames: []string{"dns.example.com", "*.example.net"}},
		{Kind: "doh", OK: false, Problem: "nothing listens", CertNames: nil},
	}}
	h.RequireNames([]string{"dns.example.com", "doh.example.net"})
	if !h.Probes[0].OK {
		t.Errorf("covered names: %+v", h.Probes[0])
	}
	h.RequireNames([]string{"dns.example.com", "doh.example.com"})
	if p := h.Probes[0]; p.OK || !strings.Contains(p.Problem, "doesn't cover doh.example.com") {
		t.Errorf("missing name: %+v", p)
	}
	if h.Probes[1].Problem != "nothing listens" {
		t.Errorf("other problem replaced: %+v", h.Probes[1])
	}
}