package blacktip

import (
	"context"
	"crypto/tls"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestParseSTS(t *testing.T) {
	good := "version: STSv1\r\nmode: enforce\r\nmx: mail.example.com\r\nmx: *.mail.example.com\r\nmax_age: 86400\r\n"
	p, err := parseSTS(good)
	if err != nil || p.Mode != "enforce" || len(p.MX) != 2 || p.MaxAge != 24*time.Hour {
		t.Fatalf("parse: %+v %v", p, err)
	}
	if p, _ := parseSTS(strings.ReplaceAll(good, "86400", "999999999999")); p == nil || p.MaxAge != 31557600*time.Second {
		t.Errorf("max_age must be capped at a year: %+v", p)
	}
	if p, err := parseSTS("version: STSv1\nmode: none\nmax_age: 60\n"); err != nil || p.Mode != "none" {
		t.Errorf("mode none needs no mx: %+v %v", p, err)
	}
	for name, bad := range map[string]string{
		"wrong version": "version: STSv2\nmode: enforce\nmx: a.b\nmax_age: 1\n",
		"unknown mode":  "version: STSv1\nmode: maybe\nmx: a.b\nmax_age: 1\n",
		"no mx":         "version: STSv1\nmode: enforce\nmax_age: 1\n",
		"no max_age":    "version: STSv1\nmode: enforce\nmx: a.b\n",
		"bad max_age":   "version: STSv1\nmode: enforce\nmx: a.b\nmax_age: soon\n",
		"empty":         "",
	} {
		if _, err := parseSTS(bad); err == nil {
			t.Errorf("accepted a policy with %s", name)
		}
	}
}

func TestSTSMXMatching(t *testing.T) {
	p := &STSPolicy{MX: []string{"mail.example.com", "*.mx.example.com"}}
	for host, want := range map[string]bool{
		"mail.example.com":      true,
		"MAIL.Example.com.":     true, // case and a trailing dot don't matter
		"a.mx.example.com":      true,
		"a.b.mx.example.com":    false, // a wildcard is exactly one label
		"mx.example.com":        false, // and never zero labels
		"evil-mail.example.com": false,
		"mail.example.com.evil": false,
		"":                      false,
	} {
		if got := p.Allows(host); got != want {
			t.Errorf("Allows(%q) = %v, want %v", host, got, want)
		}
	}
}

func TestFetchSTSHandlesTheHTTPSExchange(t *testing.T) {
	var gotHost, gotPath string
	srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotHost, gotPath = r.Host, r.URL.Path
		switch r.Header.Get("X-Case") {
		case "redirect":
			http.Redirect(w, r, "https://elsewhere.invalid/", http.StatusFound)
		case "missing":
			http.NotFound(w, r)
		default:
			w.Write([]byte("version: STSv1\nmode: testing\nmx: mx.example.test\nmax_age: 600\n"))
		}
	}))
	defer srv.Close()
	// a client that sends every request to the test server regardless of the name
	client := func(c string) *http.Client {
		return &http.Client{
			Transport: &roundTripper{c: c, rt: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
				DialContext: func(ctx context.Context, n, a string) (net.Conn, error) {
					return (&net.Dialer{}).DialContext(ctx, n, strings.TrimPrefix(srv.URL, "https://"))
				}}},
			CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
		}
	}
	p, err := fetchSTS(client(""), "example.test")
	if err != nil || p.Mode != "testing" || !p.Allows("mx.example.test") {
		t.Fatalf("fetch: %+v %v", p, err)
	}
	if gotHost != "mta-sts.example.test" || gotPath != "/.well-known/mta-sts.txt" {
		t.Errorf("fetched %s%s, want mta-sts.<domain>/.well-known/mta-sts.txt", gotHost, gotPath)
	}
	if _, err := fetchSTS(client("missing"), "example.test"); err == nil {
		t.Error("a 404 policy must be an error")
	}
	if _, err := fetchSTS(client("redirect"), "example.test"); err == nil {
		t.Error("a redirect must not be followed or accepted")
	}
}

type roundTripper struct {
	c  string
	rt http.RoundTripper
}

func (r *roundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	req = req.Clone(req.Context())
	req.Header.Set("X-Case", r.c)
	return r.rt.RoundTrip(req)
}