package server

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"regexp"
	"strings"
	"testing"

	"cc.tiredsysadmin.keeporg/internal/store"
)

// fakeIdP is the smallest thing that behaves like an OIDC provider: a
// discovery document, an authorization endpoint that redirects straight
// back, a token endpoint and a userinfo endpoint.
type fakeIdP struct {
	*httptest.Server
	subject, email, username string
	emailVerified            bool
	// tokenAuth records the client credentials the server presented.
	tokenAuth [2]string
	authCalls int
	// challenge is the PKCE challenge of the last authorization request.
	challenge string
}

func newFakeIdP(t *testing.T, subject, email, username string) *fakeIdP {
	idp := &fakeIdP{subject: subject, email: email, username: username}
	mux := http.NewServeMux()
	idp.Server = httptest.NewServer(mux)
	t.Cleanup(idp.Close)

	mux.HandleFunc("GET /.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]string{
			"issuer":                 idp.URL,
			"authorization_endpoint": idp.URL + "/authorize",
			"token_endpoint":         idp.URL + "/token",
			"userinfo_endpoint":      idp.URL + "/userinfo",
		})
	})
	mux.HandleFunc("GET /authorize", func(w http.ResponseWriter, r *http.Request) {
		idp.authCalls++
		q := r.URL.Query()
		if q.Get("code_challenge_method") != "S256" || q.Get("code_challenge") == "" {
			http.Error(w, "PKCE required", http.StatusBadRequest)
			return
		}
		idp.challenge = q.Get("code_challenge")
		back, err := url.Parse(q.Get("redirect_uri"))
		if err != nil {
			http.Error(w, "bad redirect_uri", http.StatusBadRequest)
			return
		}
		rq := back.Query()
		rq.Set("code", "authorization-code")
		rq.Set("state", q.Get("state"))
		back.RawQuery = rq.Encode()
		http.Redirect(w, r, back.String(), http.StatusFound)
	})
	mux.HandleFunc("POST /token", func(w http.ResponseWriter, r *http.Request) {
		user, pass, _ := r.BasicAuth()
		idp.tokenAuth = [2]string{user, pass}
		_ = r.ParseForm()
		if r.FormValue("code") != "authorization-code" {
			http.Error(w, "bad code", http.StatusBadRequest)
			return
		}
		if pkceChallenge(r.FormValue("code_verifier")) != idp.challenge {
			http.Error(w, "PKCE verifier does not match", http.StatusBadRequest)
			return
		}
		_ = json.NewEncoder(w).Encode(map[string]any{
			"access_token": "an-access-token", "token_type": "Bearer", "expires_in": 300,
		})
	})
	mux.HandleFunc("GET /userinfo", func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer an-access-token" {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		_ = json.NewEncoder(w).Encode(map[string]any{
			"sub": idp.subject, "email": idp.email, "preferred_username": idp.username,
			"email_verified": idp.emailVerified,
		})
	})
	return idp
}

// withOIDC configures a harness against this provider.
func (idp *fakeIdP) withOIDC(autoCreate bool, defaultGroup string) func(*Config) {
	return func(c *Config) {
		c.OIDC = OIDCConfig{
			Issuer: idp.URL, ClientID: "keeporg", ClientSecret: "a-client-secret",
			AutoCreate: autoCreate, DefaultGroupID: defaultGroup,
		}
	}
}

func TestOIDCSignsInToTheWebInterface(t *testing.T) {
	idp := newFakeIdP(t, "subject-1", "sso@example.com", "ssouser")
	h := newHarnessWith(t, idp.withOIDC(true, "everyone"))
	b := newBrowser(t, h)

	// The sign-in page offers it, and following that link is the whole flow.
	status, page := b.get("/web")
	if status != http.StatusOK || !strings.Contains(page, "/oidc/login") {
		t.Fatalf("the sign-in page does not offer single sign-on:\n%s", page)
	}
	status, page = b.get("/oidc/login")
	if status != http.StatusOK {
		t.Fatalf("oidc login: %d\n%s", status, page)
	}
	if !strings.Contains(page, "Your vaults") {
		t.Fatalf("single sign-on did not land on the vault list:\n%s", page)
	}
	if b.cookie() == "" {
		t.Fatal("no session cookie after signing in with single sign-on")
	}
	if idp.authCalls != 1 {
		t.Fatalf("the provider's authorize endpoint was called %d times", idp.authCalls)
	}
	if idp.tokenAuth != [2]string{"keeporg", "a-client-secret"} {
		t.Fatalf("the server authenticated to the token endpoint as %v", idp.tokenAuth)
	}

	// The account was provisioned, linked to the subject, and put in the
	// configured group.
	var linked *store.User
	h.srv.Store().Read(func(st *store.State) {
		linked = st.UserByOIDC(idp.URL, "subject-1")
	})
	if linked == nil {
		t.Fatal("no account was linked to the identity")
	}
	if linked.Username != "ssouser" || linked.Email != "sso@example.com" {
		t.Fatalf("provisioned account = %+v", linked)
	}
	var inGroup bool
	h.srv.Store().Read(func(st *store.State) {
		for _, m := range st.Memberships {
			if m.UserID == linked.ID && m.GroupID == "everyone" {
				inGroup = true
			}
		}
	})
	if !inGroup {
		t.Fatal("the new account did not join the default group")
	}
}

var noncePattern = regexp.MustCompile(`name="nonce" value="([^"]+)"`)

// startDeviceApproval starts a device login and takes a browser through
// single sign-on for it, up to the confirmation page.
func startDeviceApproval(t *testing.T, h *harness, b *browser) (start map[string]any, nonce string) {
	t.Helper()
	status, start := h.call("POST", "/api/v1/auth/device/start", "",
		map[string]string{"label": "phone", "return_url": "keeporg://device-approved"})
	if status != http.StatusOK {
		t.Fatalf("device start: %d %v", status, start)
	}
	status, page := b.get("/oidc/login?code=" + url.QueryEscape(start["user_code"].(string)))
	if status != http.StatusOK || !strings.Contains(page, "Let a device in?") ||
		!strings.Contains(page, start["user_code"].(string)) {
		t.Fatalf("no confirmation page after single sign-on: %d\n%s", status, page)
	}
	m := noncePattern.FindStringSubmatch(page)
	if m == nil {
		t.Fatalf("no nonce on the confirmation page:\n%s", page)
	}
	return start, m[1]
}

func TestOIDCApprovesAWaitingClient(t *testing.T) {
	idp := newFakeIdP(t, "subject-2", "dev@example.com", "devuser")
	h := newHarnessWith(t, idp.withOIDC(true, "everyone"))
	b := newBrowser(t, h)
	start, nonce := startDeviceApproval(t, h, b)
	poll := func() (int, map[string]any) {
		return h.call("POST", "/api/v1/auth/device/poll", "",
			map[string]string{"device_code": start["device_code"].(string)})
	}

	// Signing in at the provider is not yet an approval: one click on a
	// link must never let somebody else's device in.
	if status, _ := poll(); status != http.StatusAccepted {
		t.Fatalf("the device was approved before the person said yes: %d", status)
	}
	_, page := b.post("/oidc/approve", url.Values{"nonce": {nonce}, "decision": {"approve"}})
	if !strings.Contains(page, "Device approved") || !strings.Contains(page, `href="keeporg://device-approved"`) {
		t.Fatalf("approval page:\n%s", page)
	}
	// A device approval must not also hand the browser a session cookie.
	if b.cookie() != "" {
		t.Fatal("approving a device set a web session cookie")
	}
	status, body := poll()
	if status != http.StatusOK || body["token"] == "" {
		t.Fatalf("poll after approval: %d %v", status, body)
	}
	if u := body["user"].(map[string]any); u["username"] != "devuser" {
		t.Fatalf("the session belongs to %v", u["username"])
	}
}

func TestOIDCDeviceApprovalCanBeRefused(t *testing.T) {
	idp := newFakeIdP(t, "subject-6", "", "devuser")
	h := newHarnessWith(t, idp.withOIDC(true, ""))
	b := newBrowser(t, h)
	start, nonce := startDeviceApproval(t, h, b)
	_, page := b.post("/oidc/approve", url.Values{"nonce": {nonce}, "decision": {"deny"}})
	if !strings.Contains(page, "not signed in") {
		t.Fatalf("refusal page:\n%s", page)
	}
	if status, _ := h.call("POST", "/api/v1/auth/device/poll", "",
		map[string]string{"device_code": start["device_code"].(string)}); status != http.StatusGone {
		t.Fatalf("a refused device can still poll: %d", status)
	}
}

// The attack: someone starts a device login and gets the victim's browser to
// the provider with that device's code. The victim's browser ends on the
// confirmation page; the approval itself cannot be posted from anywhere
// else, because it needs this browser's cookie.
func TestOIDCApprovalCannotBePostedFromAnotherBrowser(t *testing.T) {
	idp := newFakeIdP(t, "subject-7", "", "devuser")
	h := newHarnessWith(t, idp.withOIDC(true, ""))
	victim := newBrowser(t, h)
	start, nonce := startDeviceApproval(t, h, victim)

	attacker := newBrowser(t, h)
	_, page := attacker.post("/oidc/approve", url.Values{"nonce": {nonce}, "decision": {"approve"}})
	if !strings.Contains(page, "not started in this browser") {
		t.Fatalf("an approval was accepted from another browser:\n%s", page)
	}
	if status, _ := h.call("POST", "/api/v1/auth/device/poll", "",
		map[string]string{"device_code": start["device_code"].(string)}); status != http.StatusAccepted {
		t.Fatalf("the device got in: %d", status)
	}
}

func TestOIDCNeverLinksByUsername(t *testing.T) {
	idp := newFakeIdP(t, "subject-3", "", "admin")
	h := newHarnessWith(t, idp.withOIDC(true, ""))
	b := newBrowser(t, h)
	b.get("/oidc/login")
	h.srv.Store().Read(func(st *store.State) {
		if st.Users["admin"].OIDCSubject != "" {
			t.Fatal("an identity calling itself admin claimed the admin account")
		}
		if u := st.UserByOIDC(idp.URL, "subject-3"); u == nil || u.Username == "admin" || u.ServerAdmin {
			t.Fatalf("the identity got %+v", u)
		}
	})
}

func TestOIDCLinksAVerifiedEmailOnlyWhenAllowed(t *testing.T) {
	for _, c := range []struct {
		name            string
		allow, verified bool
		wantLinkedToBob bool
	}{
		{"off", false, true, false},
		{"unverified", true, false, false},
		{"verified", true, true, true},
	} {
		t.Run(c.name, func(t *testing.T) {
			idp := newFakeIdP(t, "subject-8", "bob@example.com", "someone-else")
			idp.emailVerified = c.verified
			h := newHarnessWith(t, func(cfg *Config) {
				idp.withOIDC(false, "")(cfg)
				cfg.OIDC.LinkVerifiedEmail = c.allow
			})
			_ = h.srv.Store().Update(func(st *store.State) error {
				st.Users["bob"] = &store.User{ID: "bob", Username: "bob", Email: "bob@example.com", PassHash: "x"}
				return nil
			})
			newBrowser(t, h).get("/oidc/login")
			var linked bool
			h.srv.Store().Read(func(st *store.State) { linked = st.Users["bob"].OIDCSubject == "subject-8" })
			if linked != c.wantLinkedToBob {
				t.Fatalf("linked = %v, want %v", linked, c.wantLinkedToBob)
			}
		})
	}
}

func TestOIDCIsLinkedFromTheAccountPage(t *testing.T) {
	idp := newFakeIdP(t, "subject-9", "", "anything")
	h := newHarnessWith(t, idp.withOIDC(false, ""))
	b := newBrowser(t, h)
	csrf := b.signIn()

	// Without the session's CSRF token, a link elsewhere cannot do it.
	b.get("/oidc/login?link=1")
	_, page := b.get("/web/account")
	if !strings.Contains(page, "Link single sign-on") {
		t.Fatalf("the account page does not offer linking:\n%s", page)
	}
	_, page = b.get("/oidc/login?link=1&csrf=" + url.QueryEscape(csrf))
	if !strings.Contains(page, "single sign-on") || strings.Contains(page, "Link single sign-on") {
		t.Fatalf("not linked:\n%s", page)
	}
	h.srv.Store().Read(func(st *store.State) {
		if st.Users["admin"].OIDCSubject != "subject-9" {
			t.Fatalf("admin = %+v", st.Users["admin"])
		}
	})
}

func TestOIDCRefusesAnUnknownIdentityWhenProvisioningIsOff(t *testing.T) {
	idp := newFakeIdP(t, "subject-4", "stranger@example.com", "stranger")
	h := newHarnessWith(t, idp.withOIDC(false, ""))
	b := newBrowser(t, h)

	_, page := b.get("/oidc/login")
	if !strings.Contains(page, "no account is linked") {
		t.Fatalf("an unknown identity was let in:\n%s", page)
	}
	if b.cookie() != "" {
		t.Fatal("a session cookie was issued to an unknown identity")
	}
}

func TestOIDCCallbackRejectsAnUnknownState(t *testing.T) {
	idp := newFakeIdP(t, "subject-5", "", "someone")
	h := newHarnessWith(t, idp.withOIDC(true, ""))
	b := newBrowser(t, h)

	_, page := b.get("/oidc/callback?code=authorization-code&state=not-a-real-state")
	if !strings.Contains(page, "not started in this browser") {
		t.Fatalf("a forged callback was accepted:\n%s", page)
	}
	if b.cookie() != "" {
		t.Fatal("a forged callback issued a session cookie")
	}
}

func TestSignInPageHidesSingleSignOnWhenItIsNotConfigured(t *testing.T) {
	b := newBrowser(t, newHarness(t))
	_, page := b.get("/web")
	if strings.Contains(page, "oidc/login") {
		t.Fatal("single sign-on is offered although no provider is configured")
	}
}