st-pierre/auth.go

package main

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"strings"
	"time"

	"go.uber.org/zap"
	"golang.org/x/crypto/bcrypt"
)

var sessionTokenPrefix = "stp_"

type userRow struct {
	ID           int64  `json:"id"`
	Username     string `json:"username"`
	PasswordHash string `json:"-"`
	CreatedAt    string `json:"created_at"`
	Active       bool   `json:"active"`
}

type sessionRow struct {
	ID         int64  `json:"id"`
	UserID     int64  `json:"user_id"`
	Label      string `json:"label"`
	TokenHash  string `json:"-"`
	CreatedAt  string `json:"created_at"`
	LastUsedAt string `json:"last_used_at,omitempty"`
	Active     bool   `json:"active"`
}

type loginRequest struct {
	Username string `json:"username"`
	Password string `json:"password"`
}

type createSessionResponse struct {
	ID        int64  `json:"id"`
	Label     string `json:"label"`
	Token     string `json:"token"`
	CreatedAt string `json:"created_at"`
}

func hashPassword(password string) (string, error) {
	bytes, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		return "", err
	}
	return string(bytes), nil
}

func checkPassword(hashed, password string) bool {
	return bcrypt.CompareHashAndPassword([]byte(hashed), []byte(password)) == nil
}

func generateToken() (string, error) {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return sessionTokenPrefix + hex.EncodeToString(b), nil
}

func hashToken(token string) string {
	h := sha256.Sum256([]byte(token))
	return hex.EncodeToString(h[:])
}

func (s *sqliteStore) countUsers(ctx context.Context) (int, error) {
	var n int
	err := s.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM users`).Scan(&n)
	return n, err
}

func (s *sqliteStore) createUser(ctx context.Context, username, password string) (*userRow, error) {
	h, err := hashPassword(password)
	if err != nil {
		return nil, fmt.Errorf("hash password: %w", err)
	}

	result, err := s.db.ExecContext(ctx,
		`INSERT INTO users (username, password_hash) VALUES (?, ?)`,
		username, h,
	)
	if err != nil {
		return nil, fmt.Errorf("insert user: %w", err)
	}

	id, _ := result.LastInsertId()
	return s.getUserByID(ctx, id)
}

func (s *sqliteStore) getUserByUsername(ctx context.Context, username string) (*userRow, error) {
	row := s.db.QueryRowContext(ctx,
		`SELECT id, username, password_hash, created_at, active FROM users WHERE username = ?`, username,
	)

	var u userRow
	var active int
	err := row.Scan(&u.ID, &u.Username, &u.PasswordHash, &u.CreatedAt, &active)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("query user: %w", err)
	}
	u.Active = active == 1
	return &u, nil
}

func (s *sqliteStore) getUserByID(ctx context.Context, id int64) (*userRow, error) {
	row := s.db.QueryRowContext(ctx,
		`SELECT id, username, password_hash, created_at, active FROM users WHERE id = ?`, id,
	)

	var u userRow
	var active int
	err := row.Scan(&u.ID, &u.Username, &u.PasswordHash, &u.CreatedAt, &active)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("query user: %w", err)
	}
	u.Active = active == 1
	return &u, nil
}

func (s *sqliteStore) listUsers(ctx context.Context) ([]userRow, error) {
	rows, err := s.db.QueryContext(ctx,
		`SELECT id, username, created_at, active FROM users ORDER BY id`,
	)
	if err != nil {
		return nil, fmt.Errorf("list users: %w", err)
	}
	defer rows.Close()

	var users []userRow
	for rows.Next() {
		var u userRow
		var active int
		if err := rows.Scan(&u.ID, &u.Username, &u.CreatedAt, &active); err != nil {
			return nil, fmt.Errorf("scan user: %w", err)
		}
		u.Active = active == 1
		users = append(users, u)
	}
	return users, rows.Err()
}

func (s *sqliteStore) deactivateUser(ctx context.Context, id int64) error {
	_, err := s.db.ExecContext(ctx, `UPDATE users SET active = 0 WHERE id = ?`, id)
	return err
}

func (s *sqliteStore) changePassword(ctx context.Context, userID int64, newHash string) error {
	_, err := s.db.ExecContext(ctx,
		`UPDATE users SET password_hash = ? WHERE id = ?`, newHash, userID,
	)
	return err
}

func (s *sqliteStore) createSession(ctx context.Context, userID int64, token, label string) (*sessionRow, error) {
	h := hashToken(token)
	result, err := s.db.ExecContext(ctx,
		`INSERT INTO sessions (user_id, label, token_hash) VALUES (?, ?, ?)`,
		userID, label, h,
	)
	if err != nil {
		return nil, fmt.Errorf("insert session: %w", err)
	}

	id, _ := result.LastInsertId()
	row := s.db.QueryRowContext(ctx,
		`SELECT id, user_id, label, created_at, active FROM sessions WHERE id = ?`, id,
	)

	var sr sessionRow
	var active int
	err = row.Scan(&sr.ID, &sr.UserID, &sr.Label, &sr.CreatedAt, &active)
	if err != nil {
		return nil, fmt.Errorf("query session: %w", err)
	}
	sr.Active = active == 1
	return &sr, nil
}

func (s *sqliteStore) getSessionByToken(ctx context.Context, token string) (*sessionRow, *userRow, error) {
	h := hashToken(token)
	row := s.db.QueryRowContext(ctx,
		`SELECT s.id, s.user_id, s.label, s.created_at, s.last_used_at, s.active,
		        u.id, u.username, u.created_at, u.active
		 FROM sessions s
		 JOIN users u ON u.id = s.user_id
		 WHERE s.token_hash = ?`, h,
	)

	var sr sessionRow
	var ur userRow
	var sActive, uActive int
	var lastUsed sql.NullString

	err := row.Scan(&sr.ID, &sr.UserID, &sr.Label, &sr.CreatedAt, &lastUsed, &sActive,
		&ur.ID, &ur.Username, &ur.CreatedAt, &uActive)
	if err == sql.ErrNoRows {
		return nil, nil, nil
	}
	if err != nil {
		return nil, nil, fmt.Errorf("query session: %w", err)
	}

	sr.Active = sActive == 1
	if lastUsed.Valid {
		sr.LastUsedAt = lastUsed.String
	}
	ur.Active = uActive == 1
	return &sr, &ur, nil
}

func (s *sqliteStore) listSessions(ctx context.Context, userID int64) ([]sessionRow, error) {
	rows, err := s.db.QueryContext(ctx,
		`SELECT id, user_id, label, created_at, last_used_at, active
		 FROM sessions WHERE user_id = ? ORDER BY id`, userID,
	)
	if err != nil {
		return nil, fmt.Errorf("list sessions: %w", err)
	}
	defer rows.Close()

	var sessions []sessionRow
	for rows.Next() {
		var sr sessionRow
		var active int
		var lastUsed sql.NullString
		if err := rows.Scan(&sr.ID, &sr.UserID, &sr.Label, &sr.CreatedAt, &lastUsed, &active); err != nil {
			return nil, fmt.Errorf("scan session: %w", err)
		}
		sr.Active = active == 1
		if lastUsed.Valid {
			sr.LastUsedAt = lastUsed.String
		}
		sessions = append(sessions, sr)
	}
	return sessions, rows.Err()
}

func (s *sqliteStore) deactivateSession(ctx context.Context, id, userID int64) error {
	_, err := s.db.ExecContext(ctx,
		`UPDATE sessions SET active = 0 WHERE id = ? AND user_id = ?`, id, userID,
	)
	return err
}

func (s *sqliteStore) updateSessionLastUsed(ctx context.Context, id int64) {
	s.db.ExecContext(ctx,
		`UPDATE sessions SET last_used_at = datetime('now') WHERE id = ?`, id,
	)
}

type contextKey string

const userContextKey contextKey = "user"

func getUserFromContext(r *http.Request) *userRow {
	if u, ok := r.Context().Value(userContextKey).(*userRow); ok {
		return u
	}
	return nil
}

func authMiddleware(store *sqliteStore, logger *zap.Logger, next http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodOptions {
			next.ServeHTTP(w, r)
			return
		}

		token := ""

		cookie, err := r.Cookie("st_pierre_session")
		if err == nil && cookie.Value != "" {
			token = cookie.Value
		}

		if token == "" {
			authHeader := r.Header.Get("Authorization")
			if strings.HasPrefix(authHeader, "Bearer ") {
				token = strings.TrimPrefix(authHeader, "Bearer ")
			}
		}

		if token == "" {
			writeUnauthorized(w)
			return
		}

		session, user, err := store.getSessionByToken(r.Context(), token)
		if err != nil {
			logger.Error("auth lookup failed", zap.Error(err))
			writeUnauthorized(w)
			return
		}
		if session == nil || !session.Active || !user.Active {
			writeUnauthorized(w)
			return
		}

		store.updateSessionLastUsed(r.Context(), session.ID)

		ctx := context.WithValue(r.Context(), userContextKey, user)
		next.ServeHTTP(w, r.WithContext(ctx))
	}
}

func writeUnauthorized(w http.ResponseWriter) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusUnauthorized)
	w.Write([]byte(`{"error":"unauthorized"}`))
}

func (s *sqliteStore) bootstrapAdmin(ctx context.Context, logger *zap.Logger) {
	n, err := s.countUsers(ctx)
	if err != nil {
		logger.Fatal("failed to check users table", zap.Error(err))
	}
	if n > 0 {
		return
	}

	username := getenvDefault("ADMIN_USERNAME", "admin")
	password := os.Getenv("ADMIN_PASSWORD")
	if password == "" {
		password = generateRandomPassword()
		logger.Warn("ADMIN_PASSWORD not set, generated random password",
			zap.String("password", password))
	}

	admin, err := s.createUser(ctx, username, password)
	if err != nil {
		logger.Fatal("failed to create admin user", zap.Error(err))
	}
	logger.Info("created admin user", zap.String("username", admin.Username))

	mcpToken := os.Getenv("MCP_TOKEN")
	if mcpToken != "" {
		tokenHash := hashToken(mcpToken)
		_, err := s.db.ExecContext(ctx,
			`INSERT INTO sessions (user_id, label, token_hash) VALUES (?, ?, ?)`,
			admin.ID, "MCP_TOKEN bootstrap", tokenHash,
		)
		if err != nil {
			logger.Fatal("failed to bootstrap session from MCP_TOKEN", zap.Error(err))
		}
		logger.Info("bootstrapped session from MCP_TOKEN")
	}
}

func generateRandomPassword() string {
	b := make([]byte, 16)
	rand.Read(b)
	return hex.EncodeToString(b)
}

func getenvDefault(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}

func loginHandler(store *sqliteStore, logger *zap.Logger) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
			return
		}

		var req loginRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid request body"})
			return
		}

		if req.Username == "" || req.Password == "" {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "username and password required"})
			return
		}

		user, err := store.getUserByUsername(r.Context(), req.Username)
		if err != nil {
			logger.Error("login lookup failed", zap.Error(err))
			writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal error"})
			return
		}
		if user == nil || !user.Active || !checkPassword(user.PasswordHash, req.Password) {
			writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "invalid credentials"})
			return
		}

		token, err := generateToken()
		if err != nil {
			logger.Error("generate token failed", zap.Error(err))
			writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal error"})
			return
		}

		session, err := store.createSession(r.Context(), user.ID, token, "UI login")
		if err != nil {
			logger.Error("create session failed", zap.Error(err))
			writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal error"})
			return
		}

		http.SetCookie(w, &http.Cookie{
			Name:     "st_pierre_session",
			Value:    token,
			Path:     "/",
			HttpOnly: true,
			SameSite: http.SameSiteLaxMode,
			Expires:  time.Now().Add(365 * 24 * time.Hour),
		})

		writeJSON(w, http.StatusOK, map[string]any{
			"session": session,
			"user": map[string]any{
				"id":       user.ID,
				"username": user.Username,
			},
		})
	}
}

func logoutHandler(store *sqliteStore) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
			return
		}

		cookie, err := r.Cookie("st_pierre_session")
		if err == nil && cookie.Value != "" {
			session, _, _ := store.getSessionByToken(r.Context(), cookie.Value)
			if session != nil {
				store.deactivateSession(r.Context(), session.ID, session.UserID)
			}
		}

		http.SetCookie(w, &http.Cookie{
			Name:     "st_pierre_session",
			Value:    "",
			Path:     "/",
			HttpOnly: true,
			SameSite: http.SameSiteLaxMode,
			Expires:  time.Unix(0, 0),
		})

		writeJSON(w, http.StatusOK, map[string]string{"message": "logged out"})
	}
}

func createSessionHandler(store *sqliteStore, logger *zap.Logger) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
			return
		}

		user := getUserFromContext(r)
		if user == nil {
			writeUnauthorized(w)
			return
		}

		var body struct {
			Label string `json:"label"`
		}
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid request body"})
			return
		}
		if body.Label == "" {
			body.Label = "API session"
		}

		token, err := generateToken()
		if err != nil {
			logger.Error("generate token failed", zap.Error(err))
			writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal error"})
			return
		}

		session, err := store.createSession(r.Context(), user.ID, token, body.Label)
		if err != nil {
			logger.Error("create session failed", zap.Error(err))
			writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal error"})
			return
		}

		writeJSON(w, http.StatusCreated, createSessionResponse{
			ID:        session.ID,
			Label:     session.Label,
			Token:     token,
			CreatedAt: session.CreatedAt,
		})
	}
}

func listSessionsHandler(store *sqliteStore, logger *zap.Logger) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
			return
		}

		user := getUserFromContext(r)
		if user == nil {
			writeUnauthorized(w)
			return
		}

		sessions, err := store.listSessions(r.Context(), user.ID)
		if err != nil {
			logger.Error("list sessions failed", zap.Error(err))
			writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal error"})
			return
		}

		writeJSON(w, http.StatusOK, map[string]any{"sessions": sessions})
	}
}

func deleteSessionHandler(store *sqliteStore, logger *zap.Logger) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodDelete {
			writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
			return
		}

		user := getUserFromContext(r)
		if user == nil {
			writeUnauthorized(w)
			return
		}

		id, err := parseIDFromPath(r.URL.Path, "/api/v1/sessions/")
		if err != nil {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid session id"})
			return
		}

		if err := store.deactivateSession(r.Context(), id, user.ID); err != nil {
			logger.Error("deactivate session failed", zap.Error(err))
			writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal error"})
			return
		}

		writeJSON(w, http.StatusOK, map[string]string{"message": "session revoked"})
	}
}

func createUserHandler(store *sqliteStore, logger *zap.Logger) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
			return
		}

		currentUser := getUserFromContext(r)
		if currentUser == nil {
			writeUnauthorized(w)
			return
		}

		var body struct {
			Username string `json:"username"`
			Password string `json:"password"`
		}
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid request body"})
			return
		}
		if body.Username == "" || body.Password == "" {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "username and password required"})
			return
		}
		if len(body.Password) < 8 {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "password must be at least 8 characters"})
			return
		}

		user, err := store.createUser(r.Context(), body.Username, body.Password)
		if err != nil {
			if strings.Contains(err.Error(), "UNIQUE constraint") {
				writeJSON(w, http.StatusConflict, map[string]string{"error": "username already exists"})
				return
			}
			logger.Error("create user failed", zap.Error(err))
			writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal error"})
			return
		}

		writeJSON(w, http.StatusCreated, user)
	}
}

func listUsersHandler(store *sqliteStore, logger *zap.Logger) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
			return
		}

		currentUser := getUserFromContext(r)
		if currentUser == nil {
			writeUnauthorized(w)
			return
		}

		users, err := store.listUsers(r.Context())
		if err != nil {
			logger.Error("list users failed", zap.Error(err))
			writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal error"})
			return
		}

		writeJSON(w, http.StatusOK, map[string]any{"users": users})
	}
}

func changePasswordHandler(store *sqliteStore, logger *zap.Logger) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPut {
			writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
			return
		}

		user := getUserFromContext(r)
		if user == nil {
			writeUnauthorized(w)
			return
		}

		var body struct {
			OldPassword string `json:"old_password"`
			NewPassword string `json:"new_password"`
		}
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid request body"})
			return
		}
		if body.OldPassword == "" || body.NewPassword == "" {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "old_password and new_password required"})
			return
		}
		if len(body.NewPassword) < 8 {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "new password must be at least 8 characters"})
			return
		}
		if body.OldPassword == body.NewPassword {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "new password must differ from old"})
			return
		}

		fullUser, err := store.getUserByID(r.Context(), user.ID)
		if err != nil || fullUser == nil {
			logger.Error("get user for password change", zap.Error(err))
			writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal error"})
			return
		}

		if !checkPassword(fullUser.PasswordHash, body.OldPassword) {
			writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "incorrect old password"})
			return
		}

		newHash, err := hashPassword(body.NewPassword)
		if err != nil {
			logger.Error("hash new password", zap.Error(err))
			writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal error"})
			return
		}

		if err := store.changePassword(r.Context(), user.ID, newHash); err != nil {
			logger.Error("change password", zap.Error(err))
			writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal error"})
			return
		}

		writeJSON(w, http.StatusOK, map[string]string{"message": "password changed"})
	}
}

func parseIDFromPath(path, prefix string) (int64, error) {
	rest := strings.TrimPrefix(path, prefix)
	rest = strings.TrimSuffix(rest, "/")
	if rest == "" {
		return 0, fmt.Errorf("empty id")
	}
	var id int64
	_, err := fmt.Sscanf(rest, "%d", &id)
	return id, err
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(v)
}

func getUserHandler(store *sqliteStore, logger *zap.Logger) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
			return
		}

		user := getUserFromContext(r)
		if user == nil {
			writeUnauthorized(w)
			return
		}

		writeJSON(w, http.StatusOK, user)
	}
}