package main

import (
	"maps"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"
)

// The RFC's own vectors, which is the only way to know an authenticator will
// agree.
func TestTOTPAgainstRFC6238(t *testing.T) {
	// The RFC's SHA-1 key is the ASCII "12345678901234567890".
	secret := knownSecret
	cases := map[int64]string{
		59:         "287082",
		1111111109: "081804",
		1111111111: "050471",
		1234567890: "005924",
		2000000000: "279037",
	}
	for at, want := range cases {
		if got := totpCode(secret, time.Unix(at, 0)); got != want {
			t.Errorf("at %d: %s, wanted %s", at, got, want)
		}
	}
	if totpCode("not base32 !!", time.Now()) != "" {
		t.Error("an unreadable secret produced a code")
	}
}

func TestTOTPSkew(t *testing.T) {
	secret := totpSecret()
	now := time.Now()
	if totpStep(secret, totpCode(secret, now), now) < 0 {
		t.Error("the current code was refused")
	}
	if totpStep(secret, totpCode(secret, now.Add(-totpPeriod)), now) < 0 {
		t.Error("the previous code was refused")
	}
	if totpStep(secret, totpCode(secret, now.Add(-5*totpPeriod)), now) >= 0 {
		t.Error("a code from two minutes ago was accepted")
	}
	if totpStep(secret, "000000 ", now) >= 0 && totpCode(secret, now) != "000000" {
		t.Error("a wrong code was accepted")
	}
	// The step reported is the one the code was made for, which is what stops
	// the previous code once the current one has been used.
	if got, want := totpStep(secret, totpCode(secret, now.Add(-totpPeriod)), now), now.Unix()/30-1; got != want {
		t.Errorf("the previous code reports step %d, wanted %d", got, want)
	}
}

func TestTOTPLoginAndRecovery(t *testing.T) {
	f := testForge(t)
	addUser(t, f, "alice", "hunter2hunter2")
	h := f.routes()

	session := signIn(t, f, h, "alice")
	token := f.csrf(withCookies("/", session))

	// Enrolling needs a code that proves the authenticator has the secret.
	page := get(h, "/settings/totp", session).Body.String()
	secret := between(page, `name="secret" value="`, `"`)
	if secret == "" {
		t.Fatal("no proposed secret on the enrolment page")
	}
	// Enrolling costs the password as well as a code.
	bare := url.Values{"csrf": {token}, "secret": {secret}, "code": {totpCode(secret, time.Now())}}
	if w := post(h, "/settings/totp", bare, session); w.Code != http.StatusUnauthorized {
		t.Errorf("enrolling with no password: status %d", w.Code)
	}
	wrong := url.Values{"csrf": {token}, "secret": {secret}, "code": {"000000"}, "current": {"hunter2hunter2"}}
	retry := post(h, "/settings/totp", wrong, session)
	if retry.Code != http.StatusBadRequest {
		t.Errorf("enrolling with a wrong code: status %d", retry.Code)
	}
	// The page says to try the next code, which is only possible if the secret
	// is still the one already typed into the authenticator.
	if got := between(retry.Body.String(), `name="secret" value="`, `"`); got != secret {
		t.Errorf("a wrong code rotated the secret to %q", got)
	}
	w := post(h, "/settings/totp", url.Values{
		"csrf": {token}, "secret": {secret}, "code": {totpCode(secret, time.Now())},
		"current": {"hunter2hunter2"}}, session)
	if w.Code != http.StatusOK {
		t.Fatalf("enrolling: status %d", w.Code)
	}
	codes := recoveryFrom(w.Body.String())
	if len(codes) != recoveryCodes {
		t.Fatalf("expected %d recovery codes, got %v", recoveryCodes, codes)
	}
	if u, _ := f.store.User("alice"); u.TOTP == nil || u.TOTP.Secret != secret || u.Epoch == 0 {
		t.Errorf("enrolment did not stick, or did not bump the epoch: %+v", u.TOTP)
	}

	// The password alone now stops at the second step.
	first := post(h, "/login", url.Values{"name": {"alice"}, "password": {"hunter2hunter2"}}, nil)
	if first.Code != http.StatusOK || !strings.Contains(first.Body.String(), "Authenticator code") {
		t.Fatalf("password alone: status %d", first.Code)
	}
	if len(first.Result().Cookies()) != 0 {
		t.Error("the password alone issued a session")
	}
	pending := url.Values{
		"csrf":    {f.csrf(withCookies("/", nil))},
		"name":    {between(first.Body.String(), `name="name" value="`, `"`)},
		"expires": {between(first.Body.String(), `name="expires" value="`, `"`)},
		"t":       {between(first.Body.String(), `name="t" value="`, `"`)},
	}

	badCode := maps.Clone(pending)
	badCode.Set("code", "000000")
	if w := post(h, "/login/totp", badCode, nil); w.Code != http.StatusUnauthorized {
		t.Errorf("a wrong code: status %d", w.Code)
	}
	right := maps.Clone(pending)
	right.Set("code", freshCode(t, f, "alice", secret))
	if w := post(h, "/login/totp", right, nil); w.Code != http.StatusSeeOther || len(w.Result().Cookies()) == 0 {
		t.Errorf("the right code: status %d", w.Code)
	}
	// A code works once, however long its window has left.
	if w := post(h, "/login/totp", right, nil); w.Code != http.StatusUnauthorized || len(w.Result().Cookies()) != 0 {
		t.Errorf("the same code again: status %d", w.Code)
	}

	// A recovery code works once.
	recovery := maps.Clone(pending)
	recovery.Set("code", codes[0])
	if w := post(h, "/login/totp", recovery, nil); w.Code != http.StatusSeeOther {
		t.Errorf("a recovery code: status %d", w.Code)
	}
	if u, _ := f.store.User("alice"); len(u.TOTP.Recovery) != recoveryCodes-1 {
		t.Errorf("the recovery code was not spent: %d left", len(u.TOTP.Recovery))
	}
	if w := post(h, "/login/totp", recovery, nil); w.Code != http.StatusUnauthorized {
		t.Errorf("the same recovery code twice: status %d", w.Code)
	}
}

// The flag is not enough on its own.
func TestAdminNeedsTOTP(t *testing.T) {
	f := testForge(t)
	addUser(t, f, "alice", "hunter2hunter2")
	err := f.store.UpdateUsers("users: admin", f.store.forgeAuthor(), func(m map[string]User) error {
		u := m["alice"]
		u.Admin = true
		m["alice"] = u
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
	h := f.routes()
	session := signIn(t, f, h, "alice")

	w := get(h, "/admin", session)
	if w.Code != http.StatusSeeOther || w.Header().Get("Location") != "/settings/totp" {
		t.Fatalf("an admin with no authenticator reached /admin: %d %s", w.Code, w.Header().Get("Location"))
	}

	token := f.csrf(withCookies("/", session))
	secret := between(get(h, "/settings/totp", session).Body.String(), `name="secret" value="`, `"`)
	enrolled := post(h, "/settings/totp", url.Values{
		"csrf": {token}, "secret": {secret}, "code": {totpCode(secret, time.Now())},
		"current": {"hunter2hunter2"}}, session)
	session = enrolled.Result().Cookies()
	if get(h, "/admin", session).Code != http.StatusOK {
		t.Error("an enrolled admin still cannot reach /admin")
	}

	// And an admin may not take the second factor off again.
	token = f.csrf(withCookies("/", session))
	form := url.Values{"csrf": {token}, "current": {"hunter2hunter2"}}
	if w := post(h, "/settings/totp/remove", form, session); w.Code != http.StatusForbidden {
		t.Errorf("an admin removed their authenticator: status %d", w.Code)
	}
}

// Enrolling and reissuing recovery codes both outlive the session they were
// done from, so both cost the password.
func TestSecondFactorNeedsPassword(t *testing.T) {
	f := testForge(t)
	addUser(t, f, "victim", "hunter2hunter2")
	h := f.routes()
	stolen := signIn(t, f, h, "victim")
	token := f.csrf(withCookies("/", stolen))
	secret := knownSecret

	w := post(h, "/settings/totp", url.Values{
		"csrf": {token}, "secret": {secret}, "code": {totpCode(secret, time.Now())}}, stolen)
	if w.Code != http.StatusUnauthorized {
		t.Fatalf("a cookie alone enrolled a second factor: status %d", w.Code)
	}
	if u, _ := f.store.User("victim"); u.TOTP != nil {
		t.Fatal("the account was taken with a cookie alone")
	}
	// The retry has to offer the same secret, or the authenticator that was
	// just set up by hand no longer matches.
	if got := between(w.Body.String(), `name="secret" value="`, `"`); got != secret {
		t.Errorf("the retry offered a different secret: %q", got)
	}

	w = post(h, "/settings/totp", url.Values{
		"csrf": {token}, "secret": {secret}, "code": {totpCode(secret, time.Now())},
		"current": {"hunter2hunter2"}}, stolen)
	if w.Code != http.StatusOK {
		t.Fatalf("enrolling with the password: status %d", w.Code)
	}
	session := w.Result().Cookies()
	token = f.csrf(withCookies("/", session))

	if w := post(h, "/settings/totp/codes", url.Values{"csrf": {token}}, session); w.Code != http.StatusUnauthorized {
		t.Errorf("recovery codes regenerated without a password: status %d", w.Code)
	}
	w = post(h, "/settings/totp/codes", url.Values{"csrf": {token}, "current": {"hunter2hunter2"}}, session)
	if w.Code != http.StatusOK || len(recoveryFrom(w.Body.String())) != recoveryCodes {
		t.Errorf("regenerating with the password: status %d", w.Code)
	}
}

// An external account has no password, so the enrolled authenticator is what
// stands between a stolen cookie and losing the second factor.
func TestExternalSecondFactorNeedsACode(t *testing.T) {
	f := testForge(t)
	err := f.store.UpdateUsers("users: add sso:carol", f.store.forgeAuthor(), func(m map[string]User) error {
		m["sso:carol"] = User{Origin: originOIDC, Iss: "https://idp.example", Sub: "carol", Notify: true}
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
	addTOTP(t, f, "sso:carol")
	h := f.routes()

	// No password and no provider round trip: the cookie is all the thief has.
	stolen := httptest.NewRecorder()
	f.setSession(stolen, "sso:carol", 0)
	session := stolen.Result().Cookies()
	token := f.csrf(withCookies("/", session))

	for _, path := range []string{"/settings/totp/remove", "/settings/totp/codes"} {
		if w := post(h, path, url.Values{"csrf": {token}}, session); w.Code != http.StatusForbidden {
			t.Errorf("%s with the cookie alone: status %d", path, w.Code)
		}
		if u, _ := f.store.User("sso:carol"); u.TOTP == nil {
			t.Fatalf("%s took the second factor off with the cookie alone", path)
		}
	}

	code := url.Values{"csrf": {token}, "code": {totpCode(knownSecret, time.Now())}}
	if w := post(h, "/settings/totp/codes", code, session); w.Code != http.StatusOK {
		t.Errorf("regenerating with the code: status %d", w.Code)
	}
	code.Set("code", freshCode(t, f, "sso:carol", knownSecret))
	if w := post(h, "/settings/totp/remove", code, session); w.Code != http.StatusSeeOther {
		t.Errorf("removing with the code: status %d", w.Code)
	}
	if u, _ := f.store.User("sso:carol"); u.TOTP != nil {
		t.Error("the code was right and the authenticator is still there")
	}
}

// The way back in when both the authenticator and the recovery codes are gone.
func TestAdminClearsLostAuthenticator(t *testing.T) {
	f := testForge(t)
	addUser(t, f, "alice", "hunter2hunter2")
	addUser(t, f, "bob", "hunter2hunter2")
	addAdmin(t, f, "alice")
	addTOTP(t, f, "bob")
	h := f.routes()
	session := codeSession(t, f, h, "alice", "hunter2hunter2")
	token := f.csrf(withCookies("/", session))

	// Clearing somebody's second factor hands out access, so it asks for a code.
	bare := url.Values{"csrf": {token}, "name": {"bob"}}
	if w := post(h, "/admin/user/totp", bare, session); w.Code != http.StatusForbidden {
		t.Errorf("cleared bob's factor with no code of our own: status %d", w.Code)
	}
	if u, _ := f.store.User("bob"); u.TOTP == nil {
		t.Fatal("it was cleared anyway")
	}
	withCode := func(name string) url.Values {
		return url.Values{"csrf": {token}, "name": {name}, "code": {freshCode(t, f, "alice", knownSecret)}}
	}
	if w := post(h, "/admin/user/totp", withCode("alice"), session); w.Code != http.StatusForbidden {
		t.Errorf("an admin cleared their own factor from the admin page: status %d", w.Code)
	}
	if w := post(h, "/admin/user/totp", withCode("bob"), session); w.Code != http.StatusSeeOther {
		t.Fatalf("clearing bob: status %d", w.Code)
	}
	if u, _ := f.store.User("bob"); u.TOTP != nil || u.Epoch == 0 {
		t.Error("bob's factor was not cleared, or his sessions were left alive")
	}
	if w := post(h, "/login", url.Values{"name": {"bob"}, "password": {"hunter2hunter2"}}, nil); w.Code != http.StatusSeeOther {
		t.Errorf("bob cannot get back in: status %d", w.Code)
	}
}

// The budget is for guessing, and it is the same one the login form spends. An
// admin charged for every code that was right would lock themselves out.
func TestConfirmCodeCountsWrongCodesOnly(t *testing.T) {
	f := testForge(t)
	addUser(t, f, "alice", "hunter2hunter2")
	addUser(t, f, "bob", "hunter2hunter2")
	addAdmin(t, f, "alice")
	h := f.routes()
	session := codeSession(t, f, h, "alice", "hunter2hunter2")
	token := f.csrf(withCookies("/", session))
	// Signing in spent one of these, which is its own rule. Count from here.
	f.totpLimit = newLimiter()

	right := url.Values{"csrf": {token}, "name": {"bob"}, "admin": {"yes"}}
	for i := range 5 {
		right.Set("code", freshCode(t, f, "alice", knownSecret))
		if w := post(h, "/admin/user/admin", right, session); w.Code != http.StatusSeeOther {
			t.Fatalf("admin action %d: status %d", i, w.Code)
		}
	}
	if f.totpLimit.over(accountKey("alice"), 0) {
		t.Error("codes that were right spent the budget for guessing")
	}

	wrong := url.Values{"csrf": {token}, "name": {"bob"}, "admin": {"yes"}, "code": {"000000"}}
	for range 3 {
		if w := post(h, "/admin/user/admin", wrong, session); w.Code != http.StatusForbidden {
			t.Fatalf("a wrong code was accepted: status %d", w.Code)
		}
	}
	// The code that was just accepted is refused a second time, and counted.
	if w := post(h, "/admin/user/admin", right, session); w.Code != http.StatusForbidden {
		t.Errorf("a code used already was accepted again: status %d", w.Code)
	}
	if !f.totpLimit.over(accountKey("alice"), 3) {
		t.Error("wrong codes were not counted")
	}
	if f.totpLimit.over(accountKey("alice"), 4) {
		t.Error("more was counted than was tried")
	}
}

// The secret arrives in a form, so what is enrolled is only as strong as this
// check makes it.
func TestEnrolledSecretMustBeTheMintedSize(t *testing.T) {
	f := testForge(t)
	addUser(t, f, "alice", "hunter2hunter2")
	h := f.routes()
	session := signIn(t, f, h, "alice")
	for _, secret := range []string{"AA", "", "!!!", knownSecret + "GE"} {
		w := post(h, "/settings/totp", url.Values{
			"csrf": {f.csrf(withCookies("/", session))}, "secret": {secret},
			"code": {totpCode(secret, time.Now())}, "current": {"hunter2hunter2"},
		}, session)
		if w.Code != http.StatusBadRequest {
			t.Errorf("a secret of %d characters: status %d", len(secret), w.Code)
		}
	}
	if u, _ := f.store.User("alice"); u.TOTP != nil {
		t.Error("a secret this forge would not mint was enrolled")
	}
}

// Two logins with one recovery code, from the same read of the account: the
// write decides, not the read.
func TestRecoveryCodeIsSpentInTheWrite(t *testing.T) {
	f := testForge(t)
	addUser(t, f, "alice", "hunter2hunter2")
	show, keep := newRecoveryCodes()
	err := f.store.UpdateUser("alice", "users: enrol", f.store.forgeAuthor(), func(acct *User) error {
		acct.TOTP = &TOTP{Secret: knownSecret, Recovery: keep}
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
	stale, _ := f.store.User("alice")
	if !f.spendRecoveryCode("alice", stale, show[0]) {
		t.Fatal("a fresh recovery code was refused")
	}
	if f.spendRecoveryCode("alice", stale, show[0]) {
		t.Error("the same recovery code was accepted twice")
	}
	if f.spendRecoveryCode("alice", stale, "not-a-code") {
		t.Error("a made up recovery code was accepted")
	}
	if u, _ := f.store.User("alice"); len(u.TOTP.Recovery) != recoveryCodes-1 {
		t.Errorf("%d recovery codes left, wanted %d", len(u.TOTP.Recovery), recoveryCodes-1)
	}
}

// knownSecret is the base32 of the ASCII "12345678901234567890", the key the
// RFC's own vectors use.
const knownSecret = "GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ"

// addTOTP enrols the known authenticator on an account, so that totpCode can
// answer for it.
func addTOTP(t *testing.T, f *Forge, name string) {
	t.Helper()
	err := f.store.UpdateUsers("users: enrol "+name, f.store.forgeAuthor(), func(m map[string]User) error {
		u := m[name]
		u.TOTP = &TOTP{Secret: knownSecret, Recovery: []string{}}
		m[name] = u
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
}

// freshCode clears the step the account last accepted and returns the current
// code. A code works once, and a test cannot wait thirty seconds for the next.
func freshCode(t *testing.T, f *Forge, name, secret string) string {
	t.Helper()
	err := f.store.touchUser(name, func(acct *User) error {
		if acct.TOTP != nil {
			acct.TOTP.Used = 0
		}
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
	return totpCode(secret, time.Now())
}

// addAdmin hands out the global admin flag, with the authenticator the flag
// does nothing without.
func addAdmin(t *testing.T, f *Forge, name string) {
	t.Helper()
	addTOTP(t, f, name)
	err := f.store.UpdateUsers("users: admin "+name, f.store.forgeAuthor(), func(m map[string]User) error {
		u := m[name]
		u.Admin = true
		m[name] = u
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
}

// codeSession signs an account holding the known authenticator all the way
// in, through both steps.
func codeSession(t *testing.T, f *Forge, h http.Handler, name, password string) []*http.Cookie {
	t.Helper()
	first := post(h, "/login", url.Values{"name": {name}, "password": {password}}, nil)
	if !strings.Contains(first.Body.String(), "Authenticator code") {
		t.Fatal("expected a second step")
	}
	w := post(h, "/login/totp", url.Values{
		"csrf":    {f.csrf(withCookies("/", nil))},
		"name":    {between(first.Body.String(), `name="name" value="`, `"`)},
		"expires": {between(first.Body.String(), `name="expires" value="`, `"`)},
		"t":       {between(first.Body.String(), `name="t" value="`, `"`)},
		"code":    {totpCode(knownSecret, time.Now())},
	}, nil)
	if w.Code != http.StatusSeeOther {
		t.Fatalf("second step: %d", w.Code)
	}
	return w.Result().Cookies()
}

func between(s, start, end string) string {
	_, rest, ok := strings.Cut(s, start)
	if !ok {
		return ""
	}
	out, _, _ := strings.Cut(rest, end)
	return out
}

func recoveryFrom(page string) []string {
	block := between(page, "<pre>", "</pre>")
	var codes []string
	for _, line := range strings.Split(block, "\n") {
		if line = strings.TrimSpace(line); line != "" {
			codes = append(codes, line)
		}
	}
	return codes
}

// Removing it is the other half of enrolling: it costs the password too, and
// it takes every session with it.
func TestSecondFactorIsRemoved(t *testing.T) {
	f := testForge(t)
	addUser(t, f, "bob", "hunter2hunter2")
	addTOTP(t, f, "bob")
	h := f.routes()
	session := codeSession(t, f, h, "bob", "hunter2hunter2")
	token := f.csrf(withCookies("/", session))

	if w := post(h, "/settings/totp/remove", url.Values{"csrf": {token}}, session); w.Code != http.StatusUnauthorized {
		t.Fatalf("removing with no password: status %d", w.Code)
	}
	if u, _ := f.store.User("bob"); u.TOTP == nil {
		t.Fatal("it was removed anyway")
	}
	form := url.Values{"csrf": {token}, "current": {"hunter2hunter2"}}
	if w := post(h, "/settings/totp/remove", form, session); w.Code != http.StatusSeeOther {
		t.Fatalf("removing with the password: status %d", w.Code)
	}
	u, _ := f.store.User("bob")
	if u.TOTP != nil {
		t.Error("the authenticator is still enrolled")
	}
	if u.Epoch == 0 {
		t.Error("removing it left every other session alive")
	}
	// And the password alone gets back in.
	if w := post(h, "/login", url.Values{"name": {"bob"}, "password": {"hunter2hunter2"}}, nil); w.Code != http.StatusSeeOther {
		t.Errorf("signing in after the removal: status %d", w.Code)
	}
}