package server

// Two-factor login with TOTP (docs/permissions.md#two-factor-login).
//
// Logging in with a password is the first step for a user with TOTP: it
// returns a short-lived challenge instead of a session, and
// POST /auth/login/totp with the challenge and a code (or a recovery code)
// creates the session. A session remembers when it last proved a code;
// actions that would let a stolen password outlive the second factor
// (creating or rotating API tokens, changing users, resetting someone's
// TOTP) need one from the last stepUpWindow, or a code in the X-Jarl-TOTP
// header. API tokens and the local socket are not affected: a token is
// its own secret, and can't answer a prompt.
//
// totp.require in jarl.yaml makes TOTP compulsory for administrators or
// for everyone. A user who must use it but hasn't set it up gets a session
// that can only set it up (setupOnly).

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"

	"jarl/internal/api"
	"jarl/internal/auth"
	"jarl/internal/model"
	"jarl/internal/totp"
)

const (
	// challengeLifetime is how long the second login step may take.
	challengeLifetime = 5 * time.Minute
	// challengeTries is how many wrong codes one challenge allows.
	challengeTries = 5
	// stepUpWindow is how long a proved code counts for actions that
	// ask for a recent one.
	stepUpWindow = 5 * time.Minute
	// TOTPHeader carries a code for an action that asks for a recent one.
	TOTPHeader = "X-Jarl-TOTP"
)

// pendingLogins are logins waiting for their second step, by the hash of
// the challenge.
type pendingLogins struct {
	mu sync.Mutex
	m  map[string]*pendingLogin
}

type pendingLogin struct {
	userID  int64
	expires time.Time
	tries   int
}

func (c *pendingLogins) add(userID int64, now time.Time) string {
	id, _ := auth.NewSessionID()
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.m == nil {
		c.m = map[string]*pendingLogin{}
	}
	for k, v := range c.m {
		if now.After(v.expires) {
			delete(c.m, k)
		}
	}
	c.m[auth.HashSecret(id)] = &pendingLogin{userID: userID, expires: now.Add(challengeLifetime)}
	return id
}

func (c *pendingLogins) get(id string, now time.Time) *pendingLogin {
	c.mu.Lock()
	defer c.mu.Unlock()
	ch := c.m[auth.HashSecret(id)]
	if ch == nil || now.After(ch.expires) {
		return nil
	}
	return ch
}

// fail counts a wrong code; the challenge ends after challengeTries.
func (c *pendingLogins) fail(id string) (left int) {
	c.mu.Lock()
	defer c.mu.Unlock()
	h := auth.HashSecret(id)
	ch := c.m[h]
	if ch == nil {
		return 0
	}
	ch.tries++
	if ch.tries >= challengeTries {
		delete(c.m, h)
		return 0
	}
	return challengeTries - ch.tries
}

func (c *pendingLogins) done(id string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.m, auth.HashSecret(id))
}

// totpRequired reports whether totp.require makes TOTP compulsory for u.
// Service accounts (no password) can't log in, so it never applies to
// them.
func (s *Server) totpRequired(u *model.User) bool {
	if u == nil || u.ID == 0 || u.PasswordHash == "" {
		return false
	}
	switch s.cfg.TOTPRequire {
	case "all":
		return true
	case "admins":
		return u.Admin
	}
	return false
}

// setupOnly reports whether a session may only set up TOTP: the user
// must use it and hasn't set it up yet.
func (s *Server) setupOnly(p *auth.Principal) bool {
	return p.Method == "session" && s.totpRequired(p.User) && !p.User.HasTOTP()
}

// setupOnlyAllowed are the requests a setup-only session may make.
func setupOnlyAllowed(r *http.Request) bool {
	path := strings.TrimPrefix(r.URL.Path, api.Prefix)
	switch {
	case path == "/auth/whoami", path == "/auth/logout", path == "/me/password":
		return true
	case path == "/me/totp" || strings.HasPrefix(path, "/me/totp/"):
		return true
	}
	return false
}

var errSetupOnly = &httpError{http.StatusForbidden, "totp_setup_required", "this server requires two-factor login for your account: set it up under Account first"}

// checkCode checks a TOTP code, or with recovery set also a recovery
// code, for u, and uses it up. It returns how it was proved ("totp" or
// "recovery code").
func (s *Server) checkCode(ctx context.Context, u *model.User, code string, recovery bool) (string, error) {
	if !u.HasTOTP() {
		return "", errors.New("two-factor login is not on")
	}
	if totp.IsCode(code) {
		step, ok := totp.Verify(u.TOTPSecret, code, time.Now(), u.TOTPLastStep)
		if ok {
			if ok, err := s.store.UseTOTPStep(ctx, u.ID, step); err != nil {
				return "", err
			} else if ok {
				u.TOTPLastStep = step
				return "totp", nil
			}
		}
		return "", errors.New("wrong or already used code")
	}
	if recovery && code != "" {
		ok, err := s.store.UseRecoveryCode(ctx, u.ID, auth.HashSecret(totp.NormalizeRecovery(code)))
		if err != nil {
			return "", err
		}
		if ok {
			return "recovery code", nil
		}
		return "", errors.New("wrong or already used recovery code")
	}
	return "", errors.New("wrong code")
}

// loginChallenge answers a correct password of a user with TOTP: the
// challenge for the second step.
func (s *Server) loginChallenge(w http.ResponseWriter, u *model.User) {
	writeJSON(w, http.StatusOK, api.LoginChallenge{TOTPRequired: true, Challenge: s.pendingLogins.add(u.ID, time.Now())})
}

// loginTOTP is the second login step: POST /auth/login/totp.
func (s *Server) loginTOTP(w http.ResponseWriter, r *http.Request) error {
	var in api.LoginTOTP
	if err := decodeJSON(r, &in); err != nil {
		return err
	}
	ip := s.clientIP(r)
	now := time.Now()
	ch := s.pendingLogins.get(in.Challenge, now)
	if ch == nil {
		return &httpError{http.StatusUnauthorized, "login_expired", "the login took too long or had too many wrong codes; log in again"}
	}
	u, err := s.store.GetUser(r.Context(), ch.userID)
	if err != nil || u.Disabled {
		s.pendingLogins.done(in.Challenge)
		return &httpError{http.StatusUnauthorized, "login_failed", "your account is not available"}
	}
	userKey := "user:" + strings.ToLower(u.Username)
	who := &auth.Principal{User: u, IP: ip, Agent: r.UserAgent(), Method: "session"}
	if s.limiter.blocked("ip:"+ip, failsPerIP) || s.limiter.blocked(userKey, failsPerKey) {
		s.eng.Audit(r.Context(), who, "auth.login_blocked", "user", u.ID, u.Username, "denied", "too many failed attempts")
		return &httpError{http.StatusTooManyRequests, "rate_limited", "too many failed login attempts; wait a few minutes and try again"}
	}
	how, err := s.checkCode(r.Context(), u, in.Code, true)
	if err != nil {
		s.limiter.fail("ip:"+ip, userKey)
		left := s.pendingLogins.fail(in.Challenge)
		s.eng.Audit(r.Context(), who, "auth.totp_failed", "user", u.ID, u.Username, "failure", err.Error())
		msg := "wrong code: use the current code from your authenticator app, or a recovery code"
		if left == 0 {
			return &httpError{http.StatusUnauthorized, "login_expired", "too many wrong codes; log in again"}
		}
		return &httpError{http.StatusUnauthorized, "totp_invalid", msg}
	}
	s.pendingLogins.done(in.Challenge)
	s.limiter.reset(userKey)
	detail := "with an authenticator code"
	if how == "recovery code" {
		n, _ := s.store.RecoveryCodesLeft(r.Context(), u.ID)
		detail = fmt.Sprintf("with a recovery code (%d left)", n)
	}
	return s.startSession(w, r, u, &now, detail)
}

// startSession creates a session for u after a successful login.
// totpAt is when it proved a code (nil without TOTP).
func (s *Server) startSession(w http.ResponseWriter, r *http.Request, u *model.User, totpAt *time.Time, detail string) error {
	ip := s.clientIP(r)
	id, csrf := auth.NewSessionID()
	now := time.Now()
	se := &model.Session{IDHash: auth.HashSecret(id), UserID: u.ID, CSRFToken: csrf, CreatedAt: now, ExpiresAt: now.Add(s.cfg.SessionLifetime), LastSeenAt: now, IP: ip, UserAgent: r.UserAgent(), TOTPAt: totpAt}
	if err := s.store.CreateSession(r.Context(), se); err != nil {
		return err
	}
	s.store.TouchLogin(r.Context(), u.ID)
	http.SetCookie(w, &http.Cookie{Name: sessionCookie, Value: id, Path: "/", HttpOnly: true, Secure: s.cookieSecure(r), SameSite: http.SameSiteStrictMode, Expires: se.ExpiresAt})
	p := &auth.Principal{User: u, Method: "session", IP: ip, Agent: r.UserAgent()}
	s.eng.Audit(r.Context(), p, "auth.login", "user", u.ID, u.Username, "success", detail)
	writeJSON(w, http.StatusOK, s.whoamiOf(r.Context(), p, se))
	return nil
}

// stepUp makes sure a session of a user with TOTP proved a code within
// stepUpWindow, or proves one now from the X-Jarl-TOTP header. Tokens,
// the local socket and users without TOTP pass.
func (s *Server) stepUp(r *http.Request, p *auth.Principal) error {
	if p.Method != "session" || !p.User.HasTOTP() {
		return nil
	}
	se, _ := r.Context().Value(keySession).(*model.Session)
	if se == nil {
		return nil
	}
	now := time.Now()
	if se.TOTPAt != nil && now.Sub(*se.TOTPAt) < stepUpWindow {
		return nil
	}
	code := r.Header.Get(TOTPHeader)
	if code == "" {
		return &httpError{http.StatusForbidden, "totp_step_up", "this needs a current code from your authenticator app"}
	}
	ip := s.clientIP(r)
	userKey := "user:" + strings.ToLower(p.User.Username)
	if s.limiter.blocked(userKey, failsPerKey) {
		return &httpError{http.StatusTooManyRequests, "rate_limited", "too many wrong codes; wait a few minutes and try again"}
	}
	if _, err := s.checkCode(r.Context(), p.User, code, false); err != nil {
		s.limiter.fail("ip:"+ip, userKey)
		s.eng.Audit(r.Context(), p, "auth.totp_failed", "user", p.User.ID, p.User.Username, "failure", "step-up: "+err.Error())
		return &httpError{http.StatusForbidden, "totp_invalid", "wrong or already used code; enter the current code from your authenticator app"}
	}
	s.store.SetSessionTOTP(r.Context(), se.IDHash, now)
	se.TOTPAt = &now
	return nil
}

func (s *Server) totpStatus(ctx context.Context, u *model.User) api.TOTPStatus {
	st := api.TOTPStatus{Enabled: u.HasTOTP(), EnabledAt: u.TOTPEnabledAt, Required: s.totpRequired(u), Policy: s.cfg.TOTPRequire}
	if st.Enabled {
		st.RecoveryCodesLeft, _ = s.store.RecoveryCodesLeft(ctx, u.ID)
	}
	return st
}

// needSession refuses TOTP changes made with an API token or on the
// local socket: they belong to a person in the web UI.
func needSession(p *auth.Principal) error {
	if p.Method != "session" {
		return forbidden("two-factor login is set up in the web UI (Account), not with an API token")
	}
	return nil
}

// GET /me/totp
func (s *Server) getOwnTOTP(w http.ResponseWriter, r *http.Request, p *auth.Principal) error {
	if p.User.ID == 0 {
		return badRequest("%s on the local socket isn't a Jarl user", p.Name())
	}
	writeJSON(w, http.StatusOK, s.totpStatus(r.Context(), p.User))
	return nil
}

// POST /me/totp/setup starts setting up (or moving to a new app): a new
// secret, shown as a QR code, that becomes active once a code from it is
// confirmed. It needs the password, and with TOTP already on a recent
// code too.
func (s *Server) setupOwnTOTP(w http.ResponseWriter, r *http.Request, p *auth.Principal) error {
	if err := needSession(p); err != nil {
		return err
	}
	var in api.TOTPSetupRequest
	if err := decodeJSON(r, &in); err != nil {
		return err
	}
	u := p.User
	if u.PasswordHash != "" && !auth.CheckPassword(u.PasswordHash, in.Password) {
		s.limiter.fail("user:" + strings.ToLower(u.Username))
		s.eng.Audit(r.Context(), p, "user.totp_setup", "user", u.ID, u.Username, "failure", "wrong password")
		return &httpError{http.StatusForbidden, "wrong_password", "the password is wrong"}
	}
	if err := s.stepUp(r, p); err != nil {
		return err
	}
	secret := totp.NewSecret()
	if err := s.store.SetTOTPPending(r.Context(), u.ID, secret); err != nil {
		return err
	}
	issuer := s.cfg.TOTPIssuer
	uri := totp.URI(issuer, u.Username, secret)
	qr, err := totp.QRDataURI(uri)
	if err != nil {
		return err
	}
	writeJSON(w, http.StatusOK, api.TOTPSetup{Secret: secret, URI: uri, QR: qr, Issuer: issuer, Account: u.Username})
	return nil
}

// POST /me/totp/enable confirms the secret from setup with a code from
// the app. It returns the recovery codes (shown once) and logs out the
// user's other sessions.
func (s *Server) enableOwnTOTP(w http.ResponseWriter, r *http.Request, p *auth.Principal) error {
	if err := needSession(p); err != nil {
		return err
	}
	var in api.TOTPCode
	if err := decodeJSON(r, &in); err != nil {
		return err
	}
	u, err := s.store.GetUser(r.Context(), p.User.ID)
	if err != nil {
		return err
	}
	if u.TOTPPending == "" {
		return conflict("start the setup first (POST /me/totp/setup)")
	}
	userKey := "user:" + strings.ToLower(u.Username)
	if s.limiter.blocked(userKey, failsPerKey) {
		return &httpError{http.StatusTooManyRequests, "rate_limited", "too many wrong codes; wait a few minutes and try again"}
	}
	step, ok := totp.Verify(u.TOTPPending, in.Code, time.Now(), 0)
	if !ok {
		s.limiter.fail(userKey)
		return &httpError{http.StatusBadRequest, "totp_invalid", "the code doesn't match: check that the app shows a 6-digit code for " + s.cfg.TOTPIssuer + " and that this device's clock is right, then try the current code"}
	}
	codes := totp.NewRecoveryCodes(totp.RecoveryCodes)
	var hashes []string
	for _, c := range codes {
		hashes = append(hashes, auth.HashSecret(totp.NormalizeRecovery(c)))
	}
	replaced := u.HasTOTP()
	if err := s.store.EnableTOTP(r.Context(), u.ID, u.TOTPPending, step, hashes); err != nil {
		return err
	}
	se, _ := r.Context().Value(keySession).(*model.Session)
	keep := ""
	if se != nil {
		keep = se.IDHash
		s.store.SetSessionTOTP(r.Context(), se.IDHash, time.Now())
	}
	s.store.DeleteUserSessions(r.Context(), u.ID, keep)
	detail := "turned on"
	if replaced {
		detail = "moved to a new authenticator"
	}
	s.eng.Audit(r.Context(), p, "user.totp_enable", "user", u.ID, u.Username, "success", detail)
	writeJSON(w, http.StatusOK, api.RecoveryCodes{Codes: codes})
	return nil
}

// POST /me/totp/disable turns TOTP off, with the password and a code (or
// a recovery code). Not when totp.require applies to the user.
func (s *Server) disableOwnTOTP(w http.ResponseWriter, r *http.Request, p *auth.Principal) error {
	if err := needSession(p); err != nil {
		return err
	}
	var in api.TOTPDisable
	if err := decodeJSON(r, &in); err != nil {
		return err
	}
	u := p.User
	if !u.HasTOTP() {
		return conflict("two-factor login is not on")
	}
	if s.totpRequired(u) {
		return conflict("this server requires two-factor login for your account (totp.require in jarl.yaml); to move to a new phone, set it up again instead")
	}
	userKey := "user:" + strings.ToLower(u.Username)
	if s.limiter.blocked(userKey, failsPerKey) {
		return &httpError{http.StatusTooManyRequests, "rate_limited", "too many failed attempts; wait a few minutes and try again"}
	}
	if u.PasswordHash != "" && !auth.CheckPassword(u.PasswordHash, in.Password) {
		s.limiter.fail(userKey)
		return &httpError{http.StatusForbidden, "wrong_password", "the password is wrong"}
	}
	if _, err := s.checkCode(r.Context(), u, in.Code, true); err != nil {
		s.limiter.fail(userKey)
		return &httpError{http.StatusForbidden, "totp_invalid", err.Error()}
	}
	if err := s.store.DisableTOTP(r.Context(), u.ID); err != nil {
		return err
	}
	s.eng.Audit(r.Context(), p, "user.totp_disable", "user", u.ID, u.Username, "success", "")
	w.WriteHeader(http.StatusNoContent)
	return nil
}

// POST /me/totp/recovery-codes replaces the recovery codes (needs a
// recent code).
func (s *Server) newRecoveryCodes(w http.ResponseWriter, r *http.Request, p *auth.Principal) error {
	if err := needSession(p); err != nil {
		return err
	}
	if !p.User.HasTOTP() {
		return conflict("two-factor login is not on")
	}
	if err := s.stepUp(r, p); err != nil {
		return err
	}
	codes := totp.NewRecoveryCodes(totp.RecoveryCodes)
	var hashes []string
	for _, c := range codes {
		hashes = append(hashes, auth.HashSecret(totp.NormalizeRecovery(c)))
	}
	if err := s.store.ReplaceRecoveryCodes(r.Context(), p.User.ID, hashes); err != nil {
		return err
	}
	s.eng.Audit(r.Context(), p, "user.recovery_codes", "user", p.User.ID, p.User.Username, "success", "new recovery codes")
	writeJSON(w, http.StatusOK, api.RecoveryCodes{Codes: codes})
	return nil
}

// DELETE /users/{id}/totp: an administrator turns off someone's TOTP (a
// lost phone without recovery codes). Their sessions end; at the next
// login they log in with the password and, if totp.require applies, must
// set it up again.
func (s *Server) resetUserTOTP(w http.ResponseWriter, r *http.Request, p *auth.Principal) error {
	if err := requireAdmin(p); err != nil {
		return err
	}
	u, err := s.userByPath(r)
	if err != nil {
		return err
	}
	if err := s.stepUp(r, p); err != nil {
		return err
	}
	if !u.HasTOTP() && u.TOTPPending == "" {
		return conflict("%s doesn't use two-factor login", u.Username)
	}
	if err := s.store.DisableTOTP(r.Context(), u.ID); err != nil {
		return err
	}
	s.store.DeleteUserSessions(r.Context(), u.ID, "")
	s.eng.Audit(r.Context(), p, "user.totp_reset", "user", u.ID, u.Username, "success", "")
	u.TOTPSecret, u.TOTPEnabledAt = "", nil
	writeJSON(w, http.StatusOK, s.userOut(u))
	return nil
}

// totpOverview is the dashboard's view of two-factor login.
func (s *Server) totpOverview(ctx context.Context) *api.TOTPOverview {
	us, err := s.store.ListUsers(ctx)
	if err != nil {
		return nil
	}
	o := &api.TOTPOverview{Policy: s.cfg.TOTPRequire, Missing: []string{}}
	admins, shared := 0, ""
	for _, u := range us {
		if u.Disabled {
			continue
		}
		if s.totpRequired(u) && !u.HasTOTP() {
			o.Missing = append(o.Missing, u.Username)
		}
		if u.Admin {
			admins++
			if strings.EqualFold(u.Username, "admin") && u.PasswordHash != "" {
				shared = u.Username
			}
		}
	}
	if shared != "" && admins > 1 {
		o.SharedAdmin = shared
	}
	return o
}