// Package e2e runs hammerhead, blacktip and bamboo together in-process.
package e2e

import (
	"bytes"
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base32"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"math/big"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"git.veillax.com/Veillax/bamboo"
	"git.veillax.com/Veillax/blacktip"
	"git.veillax.com/Veillax/hammerhead"
	"git.veillax.com/Veillax/reef/model"
)

// ---------- helpers ----------

func freeUDPPort(t *testing.T) int {
	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer pc.Close()
	return pc.LocalAddr().(*net.UDPAddr).Port
}

func freePort(t *testing.T) int {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port
}

// tcpBackend answers every connection with "<tag>:<what you sent>".
func tcpBackend(t *testing.T, tag string) int {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { l.Close() })
	go func() {
		for {
			c, err := l.Accept()
			if err != nil {
				return
			}
			go func() {
				defer c.Close()
				buf := make([]byte, 4096)
				n, _ := c.Read(buf)
				fmt.Fprintf(c, "%s:%s", tag, buf[:n])
			}()
		}
	}()
	return l.Addr().(*net.TCPAddr).Port
}

func udpBackend(t *testing.T, tag string) int {
	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { pc.Close() })
	go func() {
		buf := make([]byte, 4096)
		for {
			n, addr, err := pc.ReadFrom(buf)
			if err != nil {
				return
			}
			pc.WriteTo([]byte(tag+":"+string(buf[:n])), addr)
		}
	}()
	return pc.LocalAddr().(*net.UDPAddr).Port
}

// webBackend echoes what it saw as JSON; /healthz returns healthStatus.
func webBackend(t *testing.T, name string, healthStatus int) int {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/healthz" {
			w.WriteHeader(healthStatus)
			return
		}
		json.NewEncoder(w).Encode(map[string]string{
			"backend": name, "host": r.Host, "path": r.URL.Path,
			"user": r.Header.Get("X-Hammerhead-User"), "xff": r.Header.Get("X-Forwarded-For"),
			"proto": r.Header.Get("X-Forwarded-Proto"), "cookie": r.Header.Get("Cookie"),
		})
	}))
	t.Cleanup(srv.Close)
	return srv.Listener.Addr().(*net.TCPAddr).Port
}

type client struct {
	t   *testing.T
	c   *http.Client
	u   string
	key string
}

func newClient(t *testing.T, base string) *client {
	jar, _ := cookiejar.New(nil)
	return &client{t: t, c: &http.Client{Jar: jar, Timeout: 15 * time.Second,
		CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}, u: base}
}

func (a *client) do(method, path string, in, out any) (int, string) {
	a.t.Helper()
	var body io.Reader
	if in != nil {
		b, _ := json.Marshal(in)
		body = bytes.NewReader(b)
	}
	req, _ := http.NewRequest(method, a.u+path, body)
	req.Header.Set("Content-Type", "application/json")
	if a.key != "" {
		req.Header.Set("Authorization", "Bearer "+a.key)
	}
	resp, err := a.c.Do(req)
	if err != nil {
		a.t.Fatal(err)
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(resp.Body)
	if out != nil && resp.StatusCode == 200 {
		if err := json.Unmarshal(raw, out); err != nil {
			a.t.Fatalf("%s %s: bad json %q: %v", method, path, raw, err)
		}
	}
	return resp.StatusCode, string(raw)
}

func (a *client) must(method, path string, in, out any) {
	a.t.Helper()
	if code, raw := a.do(method, path, in, out); code != 200 {
		a.t.Fatalf("%s %s -> %d %s", method, path, code, raw)
	}
}

func (a *client) expect(code int, method, path string, in any) string {
	a.t.Helper()
	got, raw := a.do(method, path, in, nil)
	if got != code {
		a.t.Fatalf("%s %s -> %d (want %d): %s", method, path, got, code, raw)
	}
	return raw
}

func tcpAsk(t *testing.T, addr, msg string) string {
	t.Helper()
	var last error
	for i := 0; i < 40; i++ {
		c, err := net.DialTimeout("tcp", addr, time.Second)
		if err == nil {
			c.SetDeadline(time.Now().Add(3 * time.Second))
			fmt.Fprint(c, msg)
			b, _ := io.ReadAll(c)
			c.Close()
			if len(b) > 0 {
				return string(b)
			}
			last = fmt.Errorf("empty reply")
		} else {
			last = err
		}
		time.Sleep(150 * time.Millisecond)
	}
	t.Fatalf("tcp %s: %v", addr, last)
	return ""
}

func udpAsk(t *testing.T, addr, msg string) string {
	t.Helper()
	var last error
	for i := 0; i < 40; i++ {
		c, err := net.Dial("udp", addr)
		if err != nil {
			t.Fatal(err)
		}
		c.SetDeadline(time.Now().Add(400 * time.Millisecond))
		c.Write([]byte(msg))
		buf := make([]byte, 4096)
		n, err := c.Read(buf)
		c.Close()
		if err == nil {
			return string(buf[:n])
		}
		last = err
	}
	t.Fatalf("udp %s: %v", addr, last)
	return ""
}

func totpNow(secret string) string {
	key, _ := base32.StdEncoding.WithPadding(base32.NoPadding).DecodeString(secret)
	var b [8]byte
	binary.BigEndian.PutUint64(b[:], uint64(time.Now().Unix()/30))
	m := hmac.New(sha1.New, key)
	m.Write(b[:])
	sum := m.Sum(nil)
	off := sum[len(sum)-1] & 0xf
	return fmt.Sprintf("%06d", (binary.BigEndian.Uint32(sum[off:off+4])&0x7fffffff)%1000000)
}

func selfSigned(t *testing.T, host string) *tls.Certificate {
	key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	tpl := &x509.Certificate{SerialNumber: big.NewInt(1), Subject: pkix.Name{CommonName: host}, DNSNames: []string{host},
		NotBefore: time.Now().Add(-time.Hour), NotAfter: time.Now().Add(time.Hour), KeyUsage: x509.KeyUsageDigitalSignature,
		ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}}
	der, err := x509.CreateCertificate(rand.Reader, tpl, tpl, &key.PublicKey, key)
	if err != nil {
		t.Fatal(err)
	}
	return &tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key}
}

// ---------- the environment ----------

type world struct {
	t       *testing.T
	ts      *httptest.Server
	hh      *hammerhead.Server
	admin   *client
	org     string
	edge    string
	edgeTok string
	site    string // remote site
	local   string // built-in local site
	mails   chan string
	mu      sync.Mutex
}

func setup(t *testing.T, tr blacktip.TraefikConfig, opts ...func(*blacktip.Config)) *world {
	w := &world{t: t, mails: make(chan string, 16)}
	var handler http.Handler
	w.ts = httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { handler.ServeHTTP(rw, r) }))
	t.Cleanup(w.ts.Close)
	hh, err := hammerhead.New(t.TempDir(), hammerhead.Options{AdminEmail: "admin@test.local", AdminPassword: "password123", DashboardURL: w.ts.URL})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(hh.Close)
	hh.Mailer = func(to, subject, body string) error { w.mails <- to + "|" + subject + "|" + body; return nil }
	handler = hh.Handler()
	w.hh = hh

	w.admin = newClient(t, w.ts.URL)
	w.admin.must("POST", "/api/login", map[string]string{"email": "admin@test.local", "password": "password123"}, nil)
	var me struct{ Orgs []struct{ ID string } }
	w.admin.must("GET", "/api/me", nil, &me)
	w.org = me.Orgs[0].ID

	tunnelPort := freeUDPPort(t)
	var edge struct{ ID, Token string }
	w.admin.must("POST", "/api/edges", map[string]any{"name": "edge1", "tunnelAddr": "127.0.0.1:" + strconv.Itoa(tunnelPort)}, &edge)
	w.edge, w.edgeTok = edge.ID, edge.Token
	var site model.Site
	w.admin.must("POST", "/api/orgs/"+w.org+"/sites", map[string]string{"name": "home"}, &site)
	w.site = site.ID
	w.refreshLocal()

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	cert := selfSigned(t, "secure.example.test")
	btCfg := blacktip.Config{Hammerhead: w.ts.URL, Token: edge.Token, TunnelAddr: "127.0.0.1:" + strconv.Itoa(tunnelPort), DataDir: t.TempDir(),
		GetCertificate: func(*tls.ClientHelloInfo) (*tls.Certificate, error) { return cert, nil }, Traefik: tr}
	for _, o := range opts {
		o(&btCfg)
	}
	bt, err := blacktip.New(btCfg)
	if err != nil {
		t.Fatal(err)
	}
	go bt.Run(ctx)
	bb, err := bamboo.New(bamboo.Config{Hammerhead: w.ts.URL, Token: site.Token, DataDir: t.TempDir()})
	if err != nil {
		t.Fatal(err)
	}
	go bb.Run(ctx)
	return w
}

func (w *world) refreshLocal() {
	var st struct{ Sites []model.Site }
	w.admin.must("GET", "/api/orgs/"+w.org+"/state", nil, &st)
	for _, s := range st.Sites {
		if s.Local {
			w.local = s.ID
		}
	}
	if w.local == "" {
		w.t.Fatal("edge did not get a built-in local site")
	}
}

func (w *world) open(m model.Mapping) model.Mapping {
	w.t.Helper()
	m.EdgeID, m.Enabled = w.edge, true
	if m.BindIP == "" {
		m.BindIP = "127.0.0.1"
	}
	var out model.Mapping
	w.admin.must("POST", "/api/orgs/"+w.org+"/mappings", map[string]any{
		"proto": m.Proto, "edgeId": m.EdgeID, "bindIp": m.BindIP, "port": m.Port, "hostname": m.Hostname, "targets": m.Targets,
		"enabled": true, "auth": m.Auth, "rules": m.Rules, "healthPath": m.HealthPath, "healthSeconds": m.HealthSeconds,
		"name": m.Name,
	}, &out)
	return out
}

func tgt(site string, port int) []model.Target {
	return []model.Target{{SiteID: site, Host: "127.0.0.1", Port: port}}
}

// ---------- tests ----------

func TestPortsMultiIPFailoverAndValidation(t *testing.T) {
	w := setup(t, blacktip.TraefikConfig{})
	remoteA, remoteB, localSvc := tcpBackend(t, "A"), tcpBackend(t, "B"), tcpBackend(t, "LOCAL")
	udpSvc := udpBackend(t, "UDP")
	port := freePort(t) // the SAME public port on different IPs

	w.open(model.Mapping{Proto: "tcp", BindIP: "127.0.0.1", Port: port, Targets: tgt(w.site, remoteA)})
	w.open(model.Mapping{Proto: "tcp", BindIP: "127.0.0.2", Port: port, Targets: tgt(w.site, remoteB)})
	w.open(model.Mapping{Proto: "tcp", BindIP: "127.0.0.3", Port: port, Targets: tgt(w.local, localSvc)})
	w.open(model.Mapping{Proto: "udp", BindIP: "127.0.0.1", Port: port, Targets: tgt(w.site, udpSvc)})
	w.open(model.Mapping{Proto: "udp", BindIP: "127.0.0.3", Port: port, Targets: tgt(w.local, udpSvc)})

	// failover: first target is dead, second is alive
	dead := freePort(t)
	fo := freePort(t)
	w.open(model.Mapping{Proto: "tcp", BindIP: "127.0.0.4", Port: fo,
		Targets: []model.Target{{SiteID: w.local, Host: "127.0.0.1", Port: dead}, {SiteID: w.site, Host: "127.0.0.1", Port: remoteA}}})

	check := func(got, want string) {
		t.Helper()
		if got != want {
			t.Errorf("got %q want %q", got, want)
		}
	}
	pa := strconv.Itoa(port)
	check(tcpAsk(t, "127.0.0.1:"+pa, "hi"), "A:hi")
	check(tcpAsk(t, "127.0.0.2:"+pa, "hi"), "B:hi")
	check(tcpAsk(t, "127.0.0.3:"+pa, "hi"), "LOCAL:hi")
	check(udpAsk(t, "127.0.0.1:"+pa, "ping"), "UDP:ping")
	check(udpAsk(t, "127.0.0.3:"+pa, "ping"), "UDP:ping")
	for i := 0; i < 4; i++ { // round robin must keep landing on the live target
		check(tcpAsk(t, "127.0.0.4:"+strconv.Itoa(fo), "x"), "A:x")
	}

	base := "/api/orgs/" + w.org + "/mappings"
	for name, m := range map[string]map[string]any{
		"duplicate":   {"proto": "tcp", "bindIp": "127.0.0.1", "port": port},
		"wildcard":    {"proto": "tcp", "bindIp": "0.0.0.0", "port": port},
		"no targets":  {"proto": "tcp", "bindIp": "127.0.0.9", "port": port, "targets": []any{}},
		"bad udp":     {"proto": "udp", "bindIp": "127.0.0.9", "port": port, "targets": append(tgt(w.site, 1), tgt(w.site, 2)...)},
		"no hostname": {"proto": "http", "bindIp": "127.0.0.9", "port": port},
	} {
		m["edgeId"] = w.edge
		if _, ok := m["targets"]; !ok {
			m["targets"] = tgt(w.site, remoteA)
		}
		w.admin.expect(400, "POST", base, m)
		_ = name
	}

	// dead target gets flagged unhealthy in the UI state
	var st struct {
		Mappings []struct {
			ID, BindIP string
			Status     model.MappingStatus
		}
	}
	deadline := time.Now().Add(8 * time.Second)
	for {
		w.admin.must("GET", "/api/orgs/"+w.org+"/state", nil, &st)
		good := false
		for _, m := range st.Mappings {
			if m.BindIP == "127.0.0.4" && len(m.Status.Targets) == 2 && !m.Status.Targets[0].Healthy && m.Status.Targets[1].Healthy {
				good = true
			}
		}
		if good {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("health state never showed the dead target: %+v", st.Mappings)
		}
		time.Sleep(250 * time.Millisecond)
	}

	// closing a port frees it
	for _, m := range st.Mappings {
		if m.BindIP == "127.0.0.2" {
			w.admin.must("DELETE", base+"/"+m.ID, nil, nil)
		}
	}
	time.Sleep(600 * time.Millisecond)
	if c, err := net.DialTimeout("tcp", "127.0.0.2:"+pa, time.Second); err == nil {
		c.Close()
		t.Error("port still open after closing it")
	}
	check(tcpAsk(t, "127.0.0.1:"+pa, "again"), "A:again")
}

type webResp struct {
	status   int
	location string
	body     map[string]string
	raw      string
}

// edgeClient talks to the edge's listener while presenting a chosen hostname.
func edgeClient(scheme string, port int) (*http.Client, func(method, host, path string) webResp) {
	jar, _ := cookiejar.New(nil)
	tr := &http.Transport{
		DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, "tcp", "127.0.0.1:"+strconv.Itoa(port))
		},
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	c := &http.Client{Jar: jar, Transport: tr, Timeout: 10 * time.Second,
		CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
	return c, func(method, host, path string) webResp {
		resp, err := c.Do(mustReq(method, scheme+"://"+host+":"+strconv.Itoa(port)+path))
		if err != nil {
			return webResp{status: -1, raw: err.Error()}
		}
		defer resp.Body.Close()
		raw, _ := io.ReadAll(resp.Body)
		out := webResp{status: resp.StatusCode, location: resp.Header.Get("Location"), raw: string(raw)}
		json.Unmarshal(raw, &out.body)
		return out
	}
}

func mustReq(method, u string) *http.Request {
	r, _ := http.NewRequest(method, u, nil)
	return r
}

// TestWebProxyAndResourceAuth runs the web suite on Blacktip's built-in proxy.
func TestWebProxyAndResourceAuth(t *testing.T) {
	webSuite(t, setup(t, blacktip.TraefikConfig{}))
}

// TestWebViaTraefik runs the identical suite with Traefik serving http/https.
// It needs a Traefik binary: set TRAEFIK_BIN to its path.
func TestWebViaTraefik(t *testing.T) {
	bin := os.Getenv("TRAEFIK_BIN")
	if bin == "" {
		t.Skip("set TRAEFIK_BIN to a traefik binary to run this test")
	}
	webSuite(t, setup(t, blacktip.TraefikConfig{Path: bin, NoACME: true}))
}

func webSuite(t *testing.T, w *world) {
	w.admin.must("POST", "/api/orgs/"+w.org+"/domains", map[string]any{"name": "example.test", "edgeId": w.edge}, nil)
	be1, be2 := webBackend(t, "one", 200), webBackend(t, "two", 500)

	port := freePort(t)
	pub := w.open(model.Mapping{Proto: "http", Port: port, Hostname: "open.example.test", Targets: tgt(w.site, be1)})

	_, get := edgeClient("http", port)
	// wait for the listener and tunnel
	var r webResp
	for i := 0; i < 60; i++ {
		if r = get("GET", "open.example.test", "/hello"); r.status == 200 {
			break
		}
		time.Sleep(150 * time.Millisecond)
	}
	if r.status != 200 || r.body["backend"] != "one" || r.body["path"] != "/hello" {
		t.Fatalf("unprotected proxy: %+v", r)
	}
	if r.body["host"] != "open.example.test:"+strconv.Itoa(port) || !strings.HasPrefix(r.body["xff"], "127.0.0.1") {
		t.Errorf("forwarded headers wrong: %+v", r.body)
	}
	if got := get("GET", "nothing.example.test", "/"); got.status != 404 {
		t.Errorf("unknown host -> %d", got.status)
	}
	_ = pub

	// ----- password + rules + share link on a protected mapping -----
	port2 := freePort(t)
	prot := w.open(model.Mapping{Proto: "http", Port: port2, Hostname: "app.example.test", Targets: tgt(w.local, be1),
		Auth:  model.Auth{Password: true, SessionHours: 1},
		Rules: []model.Rule{{Action: "deny", Match: "path", Value: "/admin"}, {Action: "allow", Match: "path", Value: "/public"}}})
	// the password itself is write-only: set it through a PUT carrying newPassword
	w.admin.must("PUT", "/api/orgs/"+w.org+"/mappings/"+prot.ID, map[string]any{
		"proto": "http", "edgeId": w.edge, "bindIp": "127.0.0.1", "port": port2, "hostname": "app.example.test", "targets": tgt(w.local, be1),
		"enabled": true, "auth": model.Auth{Password: true, SessionHours: 1}, "rules": prot.Rules, "newPassword": "hunter22",
	}, nil)

	c2, get2 := edgeClient("http", port2)
	var login webResp
	for i := 0; i < 60; i++ {
		if login = get2("GET", "app.example.test", "/dash?x=1"); login.status == 302 {
			break
		}
		time.Sleep(150 * time.Millisecond)
	}
	wantPrefix := w.ts.URL + "/auth/r/" + prot.ID + "?return="
	if login.status != 302 || !strings.HasPrefix(login.location, wantPrefix) {
		t.Fatalf("expected redirect to login, got %+v", login)
	}
	if got := get2("GET", "app.example.test", "/admin"); got.status != 403 {
		t.Errorf("deny rule -> %d", got.status)
	}
	if got := get2("GET", "app.example.test", "/public/x"); got.status != 200 || got.body["user"] != "" {
		t.Errorf("allow rule -> %+v", got)
	}
	if got := get2("POST", "app.example.test", "/dash"); got.status != 401 {
		t.Errorf("unauthenticated POST -> %d", got.status)
	}

	visitor := newClient(t, w.ts.URL)
	back := "http://app.example.test:" + strconv.Itoa(port2) + "/dash?x=1"
	visitor.expect(401, "POST", "/auth/r/"+prot.ID+"/password", map[string]string{"password": "wrong", "return": back})
	visitor.expect(400, "POST", "/auth/r/"+prot.ID+"/password", map[string]string{"password": "hunter22", "return": "http://evil.test/dash"})
	visitor.expect(400, "POST", "/auth/r/"+prot.ID+"/password", map[string]string{"password": "hunter22", "return": "http://app.example.test:1/dash"})
	var grant struct{ Redirect string }
	visitor.must("POST", "/auth/r/"+prot.ID+"/password", map[string]string{"password": "hunter22", "return": back}, &grant)
	u, _ := url.Parse(grant.Redirect)
	if u.Host != "app.example.test:"+strconv.Itoa(port2) || u.Path != "/.hammerhead/callback" {
		t.Fatalf("bad grant redirect %q", grant.Redirect)
	}
	// a forged token must be rejected
	if got := get2("GET", "app.example.test", "/.hammerhead/callback?t=abc.def&r=/"); got.status != 403 {
		t.Errorf("forged token -> %d", got.status)
	}
	cb := get2("GET", "app.example.test", u.RequestURI())
	if cb.status != 302 || cb.location != "/dash?x=1" {
		t.Fatalf("callback: %+v", cb)
	}
	got := get2("GET", "app.example.test", "/dash")
	if got.status != 200 || got.body["user"] != "password" || strings.Contains(got.body["cookie"], "hh_auth") {
		t.Fatalf("authenticated request: %+v", got)
	}
	_ = c2

	// ----- share link -----
	var share struct{ URL string }
	w.admin.must("POST", "/api/orgs/"+w.org+"/mappings/"+prot.ID+"/shares", map[string]any{"name": "friend", "maxUses": 1}, &share)
	su, _ := url.Parse(share.URL)
	if su.Host != "app.example.test" { // default port 80 isn't used in tests: rebuild the URL against the test port
		t.Logf("share url host %s", su.Host)
	}
	tok := su.Query().Get("hh_share")
	_, get3 := edgeClient("http", port2)
	sh := get3("GET", "app.example.test", "/?hh_share="+tok)
	if sh.status != 302 || !strings.Contains(sh.location, "&share="+tok) {
		t.Fatalf("share redirect: %+v", sh)
	}
	v2 := newClient(t, w.ts.URL)
	v2.must("POST", "/auth/r/"+prot.ID+"/share", map[string]string{"token": tok, "return": "http://app.example.test:" + strconv.Itoa(port2) + "/"}, &grant)
	v2.expect(401, "POST", "/auth/r/"+prot.ID+"/share", map[string]string{"token": tok, "return": back}) // single use

	// ----- SSO with roles -----
	port3 := freePort(t)
	var roles []struct {
		ID, Name string
		Admin    bool
	}
	w.admin.must("GET", "/api/orgs/"+w.org+"/roles", nil, &roles)
	var memberRole, adminRole string
	for _, r := range roles {
		if r.Name == "Member" {
			memberRole = r.ID
		}
		if r.Admin {
			adminRole = r.ID
		}
	}
	var staffRole struct{ ID string }
	w.admin.must("POST", "/api/orgs/"+w.org+"/roles", map[string]any{"name": "Staff"}, nil)
	w.admin.must("GET", "/api/orgs/"+w.org+"/roles", nil, &roles)
	for _, r := range roles {
		if r.Name == "Staff" {
			staffRole.ID = r.ID
		}
	}
	sso := w.open(model.Mapping{Proto: "http", Port: port3, Hostname: "sso.example.test", Targets: tgt(w.site, be1),
		Auth: model.Auth{SSO: true, AllowedRoles: []string{staffRole.ID}}})
	invite := func(email, role string) *client {
		var inv struct{ Link string }
		w.admin.must("POST", "/api/orgs/"+w.org+"/invites", map[string]string{"email": email, "roleId": role}, &inv)
		tok := inv.Link[strings.Index(inv.Link, "#invite=")+len("#invite="):]
		c := newClient(t, w.ts.URL)
		c.must("POST", "/api/invites/accept", map[string]string{"token": tok, "name": email, "password": "password123"}, nil)
		return c
	}
	staff := invite("staff@test.local", staffRole.ID)
	plain := invite("plain@test.local", memberRole)
	retSSO := "http://sso.example.test:" + strconv.Itoa(port3) + "/"
	plain.expect(403, "POST", "/auth/r/"+sso.ID+"/sso", map[string]string{"return": retSSO})
	staff.must("POST", "/auth/r/"+sso.ID+"/sso", map[string]string{"return": retSSO}, &grant)
	newClient(t, w.ts.URL).expect(401, "POST", "/auth/r/"+sso.ID+"/sso", map[string]string{"return": retSSO})
	_ = adminRole

	// ----- email code -----
	w.admin.must("PUT", "/api/admin/settings", map[string]any{"dashboardUrl": w.ts.URL, "smtp": map[string]any{"host": "mail.invalid", "port": 25, "from": "x@y.z"}}, nil)
	port4 := freePort(t)
	em := w.open(model.Mapping{Proto: "http", Port: port4, Hostname: "mail.example.test", Targets: tgt(w.site, be1),
		Auth: model.Auth{Email: true, AllowedMails: []string{"*@corp.test"}}})
	retEm := "http://mail.example.test:" + strconv.Itoa(port4) + "/"
	v3 := newClient(t, w.ts.URL)
	v3.must("POST", "/auth/r/"+em.ID+"/email/start", map[string]string{"email": "bob@corp.test", "return": retEm}, nil)
	var code string
	select {
	case m := <-w.mails:
		parts := strings.SplitN(m, "|", 3)
		code = strings.TrimSpace(strings.Split(strings.SplitN(parts[2], "Your code is ", 2)[1], "\n")[0])
	case <-time.After(3 * time.Second):
		t.Fatal("no email sent")
	}
	v3.expect(401, "POST", "/auth/r/"+em.ID+"/email/verify", map[string]string{"email": "bob@corp.test", "code": "000000", "return": retEm})
	v3.must("POST", "/auth/r/"+em.ID+"/email/verify", map[string]string{"email": "bob@corp.test", "code": code, "return": retEm}, &grant)
	v3.must("POST", "/auth/r/"+em.ID+"/email/start", map[string]string{"email": "eve@evil.test", "return": retEm}, nil) // same answer, no mail
	select {
	case m := <-w.mails:
		t.Errorf("mail sent to a non-allowed address: %s", m)
	case <-time.After(300 * time.Millisecond):
	}

	// ----- health check marks the 500 backend unhealthy; traffic avoids it -----
	port5 := freePort(t)
	hc := w.open(model.Mapping{Proto: "http", Port: port5, Hostname: "hc.example.test", HealthPath: "/healthz", HealthSeconds: 5,
		Targets: []model.Target{{SiteID: w.site, Host: "127.0.0.1", Port: be2}, {SiteID: w.local, Host: "127.0.0.1", Port: be1}}})
	_, get5 := edgeClient("http", port5)
	deadline := time.Now().Add(8 * time.Second)
	for {
		var st struct {
			Mappings []struct {
				ID     string
				Status model.MappingStatus
			}
		}
		w.admin.must("GET", "/api/orgs/"+w.org+"/state", nil, &st)
		done := false
		for _, m := range st.Mappings {
			if m.ID == hc.ID && len(m.Status.Targets) == 2 && !m.Status.Targets[0].Healthy && m.Status.Targets[1].Healthy {
				done = true
			}
		}
		if done {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("health check never flagged the bad backend")
		}
		time.Sleep(250 * time.Millisecond)
	}
	for i := 0; i < 5; i++ {
		if got := get5("GET", "hc.example.test", "/x"); got.status != 200 || got.body["backend"] != "one" {
			t.Fatalf("traffic reached the unhealthy backend: %+v", got)
		}
	}

	// ----- https with a certificate -----
	port6 := freePort(t)
	w.open(model.Mapping{Proto: "https", Port: port6, Hostname: "secure.example.test", Targets: tgt(w.site, be1)})
	_, get6 := edgeClient("https", port6)
	var sec webResp
	for i := 0; i < 60; i++ {
		if sec = get6("GET", "secure.example.test", "/tls"); sec.status == 200 {
			break
		}
		time.Sleep(150 * time.Millisecond)
	}
	if sec.status != 200 || sec.body["proto"] != "https" {
		t.Fatalf("https: %+v", sec)
	}

	// ----- access log and traffic counters reach hammerhead -----
	deadline = time.Now().Add(10 * time.Second)
	for {
		var logs []model.LogEntry
		w.admin.must("GET", "/api/orgs/"+w.org+"/logs?kind=access&limit=500", nil, &logs)
		sawAuth, sawDeny := false, false
		for _, l := range logs {
			if l.Mapping == prot.ID && l.Result == "auth-required" {
				sawAuth = true
			}
			if l.Mapping == prot.ID && l.Result == "denied" {
				sawDeny = true
			}
		}
		if sawAuth && sawDeny {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("access log incomplete (%d entries)", len(logs))
		}
		time.Sleep(300 * time.Millisecond)
	}
	deadline = time.Now().Add(25 * time.Second)
	for {
		var stats map[string]struct{ Conns, BytesOut int64 }
		w.admin.must("GET", "/api/orgs/"+w.org+"/stats", nil, &stats)
		if stats[pub.ID].Conns > 0 && stats[pub.ID].BytesOut > 0 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("no traffic stats: %+v", stats)
		}
		time.Sleep(500 * time.Millisecond)
	}
}

func TestAccountsOrgsAndTOTP(t *testing.T) {
	w := setup(t, blacktip.TraefikConfig{})
	base := "/api/orgs/" + w.org

	// a second org the first can't see into
	var other struct{ ID string }
	w.admin.must("POST", "/api/orgs", map[string]any{"name": "Other"}, &other)
	var roles []struct {
		ID, Name string
		Admin    bool
	}
	w.admin.must("GET", base+"/roles", nil, &roles)
	var member, adminRole string
	for _, r := range roles {
		if r.Name == "Member" {
			member = r.ID
		}
		if r.Admin {
			adminRole = r.ID
		}
	}
	invite := func(org, email, role string) *client {
		var inv struct{ Link string }
		w.admin.must("POST", "/api/orgs/"+org+"/invites", map[string]string{"email": email, "roleId": role}, &inv)
		c := newClient(t, w.ts.URL)
		c.must("POST", "/api/invites/accept", map[string]string{"token": inv.Link[strings.Index(inv.Link, "#invite=")+8:], "name": "N", "password": "password123"}, nil)
		return c
	}
	m := invite(w.org, "m@test.local", member)
	m.must("GET", base+"/state", nil, nil)
	m.expect(403, "POST", base+"/mappings", map[string]any{})
	m.expect(403, "POST", base+"/sites", map[string]string{"name": "x"})
	m.expect(403, "GET", base+"/members", nil)
	m.expect(404, "GET", "/api/orgs/"+other.ID+"/state", nil)
	m.expect(403, "POST", "/api/edges", map[string]any{"name": "e", "tunnelAddr": "1.2.3.4:5"})
	m.expect(403, "GET", "/api/admin/users", nil)

	// domains can't be claimed twice across orgs
	w.admin.must("POST", base+"/domains", map[string]any{"name": "example.test", "edgeId": w.edge}, nil)
	for _, d := range []string{"example.test", "api.example.test", "test"} {
		code, _ := w.admin.do("POST", "/api/orgs/"+other.ID+"/domains", map[string]any{"name": d}, nil)
		if code != 400 && d != "test" {
			t.Errorf("org 2 claimed %s (status %d)", d, code)
		}
	}
	// hostnames must be under the org's own domains
	w.admin.expect(400, "POST", base+"/mappings", map[string]any{"proto": "http", "edgeId": w.edge, "port": 8081,
		"hostname": "x.notmine.test", "targets": tgt(w.site, 9)})

	// the other org can't use the edge's local site (AllowLocal is off for it)
	w.admin.must("GET", "/api/orgs/"+other.ID+"/state", nil, nil)
	w.admin.expect(400, "POST", "/api/orgs/"+other.ID+"/mappings", map[string]any{"proto": "tcp", "edgeId": w.edge, "port": 25565,
		"targets": tgt(w.local, 25565)})

	// API keys act inside one org only
	var key struct{ Key string }
	w.admin.must("POST", base+"/apikeys", map[string]string{"name": "ci"}, &key)
	kc := newClient(t, w.ts.URL)
	kc.key = key.Key
	kc.must("POST", base+"/domains", map[string]any{"name": "keyed.test", "edgeId": w.edge}, nil)
	kc.expect(404, "GET", "/api/orgs/"+other.ID+"/state", nil)
	kc.expect(401, "GET", "/api/me", nil)
	bad := newClient(t, w.ts.URL)
	bad.key = "hhk_deadbeef_nope"
	bad.expect(401, "GET", base+"/state", nil)

	// removing the last admin is refused
	var members []struct{ UserID, Email string }
	w.admin.must("GET", base+"/members", nil, &members)
	for _, mm := range members {
		if mm.Email == "admin@test.local" {
			w.admin.expect(400, "PUT", base+"/members/"+mm.UserID, map[string]string{"roleId": member})
			w.admin.expect(400, "DELETE", base+"/members/"+mm.UserID, nil)
		}
	}

	// two-factor: enable, then log in with a code and with a backup code
	u := invite(w.org, "otp@test.local", adminRole)
	var setup struct{ Secret string }
	u.must("POST", "/api/me/2fa/setup", nil, &setup)
	u.expect(400, "POST", "/api/me/2fa/enable", map[string]string{"code": "000000"})
	var en struct{ BackupCodes []string }
	u.must("POST", "/api/me/2fa/enable", map[string]string{"code": totpNow(setup.Secret)}, &en)
	if len(en.BackupCodes) != 8 {
		t.Fatalf("backup codes: %v", en.BackupCodes)
	}
	fresh := newClient(t, w.ts.URL)
	var step struct {
		NeedTotp bool
		Ticket   string
	}
	fresh.must("POST", "/api/login", map[string]string{"email": "otp@test.local", "password": "password123"}, &step)
	if !step.NeedTotp {
		t.Fatal("2fa not requested")
	}
	fresh.expect(401, "GET", "/api/me", nil) // not signed in until the code is given
	fresh.expect(401, "POST", "/api/login/totp", map[string]string{"ticket": step.Ticket, "code": "123456"})
	fresh.must("POST", "/api/login/totp", map[string]string{"ticket": step.Ticket, "code": totpNow(setup.Secret)}, nil)
	fresh.must("GET", "/api/me", nil, nil)
	f2 := newClient(t, w.ts.URL)
	f2.must("POST", "/api/login", map[string]string{"email": "otp@test.local", "password": "password123"}, &step)
	f2.must("POST", "/api/login/totp", map[string]string{"ticket": step.Ticket, "code": en.BackupCodes[0]}, nil)
	f3 := newClient(t, w.ts.URL)
	f3.must("POST", "/api/login", map[string]string{"email": "otp@test.local", "password": "password123"}, &step)
	f3.expect(401, "POST", "/api/login/totp", map[string]string{"ticket": step.Ticket, "code": en.BackupCodes[0]}) // backup codes are single use

	// wrong passwords and brute force are throttled per address+email
	bf := newClient(t, w.ts.URL)
	for i := 0; i < 8; i++ {
		bf.expect(401, "POST", "/api/login", map[string]string{"email": "ghost@test.local", "password": "nope-nope"})
	}
	bf.expect(429, "POST", "/api/login", map[string]string{"email": "ghost@test.local", "password": "password123"})

	// changing the password signs out the user's other sessions
	m2 := newClient(t, w.ts.URL)
	m2.must("POST", "/api/login", map[string]string{"email": "m@test.local", "password": "password123"}, nil)
	m2.must("GET", "/api/me", nil, nil)
	m.expect(400, "POST", "/api/me/password", map[string]string{"current": "wrong", "new": "newpassword1"})
	m.expect(400, "POST", "/api/me/password", map[string]string{"current": "password123", "new": "short"})
	m.must("POST", "/api/me/password", map[string]string{"current": "password123", "new": "newpassword1"}, nil)
	m2.expect(401, "GET", "/api/me", nil)
	m.must("GET", "/api/me", nil, nil)
	newClient(t, w.ts.URL).must("POST", "/api/login", map[string]string{"email": "m@test.local", "password": "newpassword1"}, nil)
}

func TestOIDCLogin(t *testing.T) {
	w := setup(t, blacktip.TraefikConfig{})
	// a tiny identity provider
	var idp *httptest.Server
	idp = httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/.well-known/openid-configuration":
			json.NewEncoder(rw).Encode(map[string]string{"authorization_endpoint": idp.URL + "/authorize", "token_endpoint": idp.URL + "/token", "userinfo_endpoint": idp.URL + "/userinfo"})
		case "/authorize":
			q := r.URL.Query()
			if q.Get("code_challenge_method") != "S256" || q.Get("code_challenge") == "" || q.Get("client_id") != "cid" {
				http.Error(rw, "bad authorize request", 400)
				return
			}
			http.Redirect(rw, r, q.Get("redirect_uri")+"?code=thecode&state="+url.QueryEscape(q.Get("state")), 302)
		case "/token":
			r.ParseForm()
			if r.Form.Get("code") != "thecode" || r.Form.Get("client_secret") != "shh" || r.Form.Get("code_verifier") == "" {
				http.Error(rw, "bad token request", 400)
				return
			}
			json.NewEncoder(rw).Encode(map[string]string{"access_token": "at"})
		case "/userinfo":
			if r.Header.Get("Authorization") != "Bearer at" {
				http.Error(rw, "no", 401)
				return
			}
			json.NewEncoder(rw).Encode(map[string]any{"sub": "u-1", "email": "Sso.User@Corp.Test", "email_verified": true, "name": "Sso User"})
		}
	}))
	defer idp.Close()

	var roles []struct {
		ID, Name string
	}
	w.admin.must("GET", "/api/orgs/"+w.org+"/roles", nil, &roles)
	member := ""
	for _, r := range roles {
		if r.Name == "Member" {
			member = r.ID
		}
	}
	w.admin.must("POST", "/api/admin/idps", map[string]any{"name": "Corp", "issuer": idp.URL, "clientId": "cid", "clientSecret": "shh",
		"autoProvision": true, "orgId": w.org, "roleId": member, "allowedDomains": []string{"corp.test"}}, nil)

	var pub struct{ Providers []struct{ ID, Name string } }
	newClient(t, w.ts.URL).must("GET", "/api/public", nil, &pub)
	if len(pub.Providers) != 1 {
		t.Fatalf("providers: %+v", pub)
	}

	c := newClient(t, w.ts.URL)
	resp, err := c.c.Get(w.ts.URL + "/auth/oidc/" + pub.Providers[0].ID + "/start?return=/after")
	if err != nil || resp.StatusCode != 302 {
		t.Fatalf("start: %v %v", err, resp)
	}
	authz := resp.Header.Get("Location")
	resp, err = c.c.Get(authz)
	if err != nil || resp.StatusCode != 302 {
		t.Fatalf("authorize: %v %v", err, resp)
	}
	cb := resp.Header.Get("Location")
	resp, err = c.c.Get(cb)
	if err != nil || resp.StatusCode != 302 || resp.Header.Get("Location") != "/after" {
		b, _ := io.ReadAll(resp.Body)
		t.Fatalf("callback: %v %d %s", err, resp.StatusCode, b)
	}
	var me struct {
		Email string
		Orgs  []struct{ Name, Role string }
	}
	c.must("GET", "/api/me", nil, &me)
	if me.Email != "sso.user@corp.test" || len(me.Orgs) != 1 || me.Orgs[0].Role != "Member" {
		t.Fatalf("provisioned user: %+v", me)
	}
	// replaying the same state is refused
	resp, _ = c.c.Get(cb)
	if resp.StatusCode != 400 {
		t.Errorf("state replay -> %d", resp.StatusCode)
	}
}