insertr/internal/auth/auth.go

package auth

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/coreos/go-oidc/v3/oidc"
	"github.com/golang-jwt/jwt/v5"
	"golang.org/x/oauth2"
)

// UserInfo represents authenticated user information
type UserInfo struct {
	ID       string `json:"sub"`
	Email    string `json:"email,omitempty"`
	Name     string `json:"name,omitempty"`
	Provider string `json:"iss,omitempty"`
}

// AuthConfig holds authentication configuration
type AuthConfig struct {
	DevMode      bool
	Provider     string
	JWTSecret    string
	OAuthConfigs map[string]OAuthConfig
	OIDC         *OIDCConfig
}

// OAuthConfig holds OAuth provider configuration
type OAuthConfig struct {
	ClientID     string
	ClientSecret string
	RedirectURL  string
	Scopes       []string
}

// OIDCConfig holds OIDC configuration for Authentik
type OIDCConfig struct {
	Endpoint     string
	ClientID     string
	ClientSecret string
	RedirectURL  string
	Scopes       []string
}

// AuthService handles authentication operations
type AuthService struct {
	config   *AuthConfig
	provider *oidc.Provider
	oauth2   *oauth2.Config
}

// NewAuthService creates a new authentication service
func NewAuthService(config *AuthConfig) (*AuthService, error) {
	service := &AuthService{config: config}

	// Initialize OIDC provider if configured
	if config.Provider == "authentik" && config.OIDC != nil {
		if err := service.initOIDC(); err != nil {
			return nil, fmt.Errorf("failed to initialize OIDC: %w", err)
		}
	}

	return service, nil
}

// initOIDC initializes the OIDC provider and OAuth2 config
func (a *AuthService) initOIDC() error {
	ctx := context.Background()

	// Create OIDC provider
	provider, err := oidc.NewProvider(ctx, a.config.OIDC.Endpoint)
	if err != nil {
		return fmt.Errorf("failed to create OIDC provider: %w", err)
	}
	a.provider = provider

	// Default scopes if none specified
	scopes := a.config.OIDC.Scopes
	if len(scopes) == 0 {
		scopes = []string{oidc.ScopeOpenID, "profile", "email"}
	}

	// Create OAuth2 config
	a.oauth2 = &oauth2.Config{
		ClientID:     a.config.OIDC.ClientID,
		ClientSecret: a.config.OIDC.ClientSecret,
		RedirectURL:  a.config.OIDC.RedirectURL,
		Endpoint:     provider.Endpoint(),
		Scopes:       scopes,
	}

	return nil
}

// ExtractUserFromRequest extracts user information from HTTP request
func (a *AuthService) ExtractUserFromRequest(r *http.Request) (*UserInfo, error) {
	authHeader := r.Header.Get("Authorization")
	if authHeader == "" {
		return &UserInfo{ID: "anonymous"}, nil
	}

	// Parse Bearer token
	if !strings.HasPrefix(authHeader, "Bearer ") {
		return nil, fmt.Errorf("invalid authorization header format")
	}

	token := strings.TrimPrefix(authHeader, "Bearer ")

	// Handle mock tokens in development mode
	if a.config.DevMode && strings.HasPrefix(token, "mock-") {
		return a.parseMockToken(token), nil
	}

	// Parse real JWT token
	return a.parseJWT(token)
}

// parseMockToken handles mock development tokens
func (a *AuthService) parseMockToken(token string) *UserInfo {
	// Mock token format: mock-{user}-{timestamp}-{random}
	parts := strings.Split(token, "-")
	if len(parts) >= 2 {
		return &UserInfo{
			ID:       parts[1], // user part
			Email:    fmt.Sprintf("%s@localhost", parts[1]),
			Name:     strings.ToTitle(parts[1][:1]) + parts[1][1:],
			Provider: "insertr-dev",
		}
	}

	return &UserInfo{ID: "anonymous"}
}

// parseJWT parses and validates a real JWT token
func (a *AuthService) parseJWT(tokenString string) (*UserInfo, error) {
	// Use OIDC verification if available
	if a.config.Provider == "authentik" && a.provider != nil {
		return a.parseOIDCToken(tokenString)
	}

	// Fallback to HMAC verification for other tokens
	return a.parseHMACToken(tokenString)
}

// parseOIDCToken parses and validates an OIDC token using the provider's keys
func (a *AuthService) parseOIDCToken(tokenString string) (*UserInfo, error) {
	ctx := context.Background()

	// Create verifier
	verifier := a.provider.Verifier(&oidc.Config{ClientID: a.config.OIDC.ClientID})

	// Verify the token
	idToken, err := verifier.Verify(ctx, tokenString)
	if err != nil {
		return nil, fmt.Errorf("failed to verify OIDC token: %w", err)
	}

	// Extract claims
	var claims struct {
		Sub           string   `json:"sub"`
		Email         string   `json:"email"`
		EmailVerified bool     `json:"email_verified"`
		Name          string   `json:"name"`
		PreferredName string   `json:"preferred_username"`
		Groups        []string `json:"groups"`
	}

	if err := idToken.Claims(&claims); err != nil {
		return nil, fmt.Errorf("failed to extract claims: %w", err)
	}

	// Build user info
	userInfo := &UserInfo{
		ID:       claims.Sub,
		Email:    claims.Email,
		Name:     claims.Name,
		Provider: idToken.Issuer,
	}

	// Use preferred_username if name is empty
	if userInfo.Name == "" {
		userInfo.Name = claims.PreferredName
	}

	return userInfo, nil
}

// parseHMACToken parses and validates a JWT token using HMAC signing
func (a *AuthService) parseHMACToken(tokenString string) (*UserInfo, error) {
	// Parse the token
	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		// Validate signing method
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		return []byte(a.config.JWTSecret), nil
	})

	if err != nil {
		return nil, fmt.Errorf("failed to parse JWT: %w", err)
	}

	// Extract claims
	if claims, ok := token.Claims.(jwt.MapClaims); ok && token.Valid {
		userInfo := &UserInfo{}

		// Extract standard JWT claims
		if sub, ok := claims["sub"].(string); ok {
			userInfo.ID = sub
		}
		if email, ok := claims["email"].(string); ok {
			userInfo.Email = email
		}
		if name, ok := claims["name"].(string); ok {
			userInfo.Name = name
		}
		if iss, ok := claims["iss"].(string); ok {
			userInfo.Provider = iss
		}

		// Fallback to alternative claim names
		if userInfo.ID == "" {
			if userID, ok := claims["user_id"].(string); ok {
				userInfo.ID = userID
			}
		}

		// Default to anonymous if no user ID found
		if userInfo.ID == "" {
			userInfo.ID = "anonymous"
		}

		return userInfo, nil
	}

	return nil, fmt.Errorf("invalid JWT claims")
}

// CreateMockJWT creates a mock JWT token for development/testing
func (a *AuthService) CreateMockJWT(userID, email, name string) (string, error) {
	if !a.config.DevMode {
		return "", fmt.Errorf("mock JWT creation only allowed in development mode")
	}

	// Create the claims
	claims := jwt.MapClaims{
		"sub":   userID,
		"email": email,
		"name":  name,
		"iss":   "insertr-dev",
		"iat":   time.Now().Unix(),
		"exp":   time.Now().Add(24 * time.Hour).Unix(),
	}

	// Create token
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)

	// Sign and get the complete encoded token as a string
	tokenString, err := token.SignedString([]byte(a.config.JWTSecret))
	if err != nil {
		return "", fmt.Errorf("failed to sign JWT: %w", err)
	}

	return tokenString, nil
}

// ValidateToken validates a JWT token without extracting user info
func (a *AuthService) ValidateToken(tokenString string) error {
	_, err := a.parseJWT(tokenString)
	return err
}

// generateCodeVerifier generates a cryptographically secure code verifier for PKCE
func generateCodeVerifier() (string, error) {
	data := make([]byte, 32)
	if _, err := rand.Read(data); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(data), nil
}

// generateCodeChallenge generates a code challenge from the verifier using SHA256
func generateCodeChallenge(verifier string) string {
	hash := sha256.Sum256([]byte(verifier))
	return base64.RawURLEncoding.EncodeToString(hash[:])
}

// generateState generates a random state parameter for OAuth2
func generateState() string {
	data := make([]byte, 16)
	rand.Read(data)
	return base64.RawURLEncoding.EncodeToString(data)
}

// RefreshToken creates a new token with extended expiration
func (a *AuthService) RefreshToken(tokenString string) (string, error) {
	userInfo, err := a.parseJWT(tokenString)
	if err != nil {
		return "", fmt.Errorf("cannot refresh invalid token: %w", err)
	}

	// Create new token with same user info but extended expiration
	claims := jwt.MapClaims{
		"sub":   userInfo.ID,
		"email": userInfo.Email,
		"name":  userInfo.Name,
		"iss":   userInfo.Provider,
		"iat":   time.Now().Unix(),
		"exp":   time.Now().Add(24 * time.Hour).Unix(),
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString([]byte(a.config.JWTSecret))
}

// IsAuthenticated checks if the request has valid authentication
func (a *AuthService) IsAuthenticated(r *http.Request) bool {
	userInfo, err := a.ExtractUserFromRequest(r)
	return err == nil && userInfo.ID != "anonymous"
}

// RequireAuth middleware that requires authentication
func (a *AuthService) RequireAuth(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		userInfo, err := a.ExtractUserFromRequest(r)
		if err != nil || userInfo.ID == "anonymous" {
			http.Error(w, "Authentication required", http.StatusUnauthorized)
			return
		}

		// Add user info to request context for use by handlers
		ctx := r.Context()
		ctx = ContextWithUser(ctx, userInfo)
		r = r.WithContext(ctx)

		next.ServeHTTP(w, r)
	})
}

// OAuth flow handlers (to be implemented when adding real OAuth)

// HandleOAuthLogin initiates OAuth flow
func (a *AuthService) HandleOAuthLogin(w http.ResponseWriter, r *http.Request) {
	// Handle mock authentication in dev mode
	if a.config.DevMode && a.config.Provider == "mock" {
		response := map[string]interface{}{
			"message":      "Mock OAuth login",
			"redirect_url": "/auth/callback?code=mock_code&state=mock_state",
			"dev_mode":     true,
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(response)
		return
	}

	// Handle Authentik OIDC flow
	if a.config.Provider == "authentik" && a.oauth2 != nil {
		// Generate PKCE parameters
		codeVerifier, err := generateCodeVerifier()
		if err != nil {
			http.Error(w, "Failed to generate code verifier", http.StatusInternalServerError)
			return
		}

		codeChallenge := generateCodeChallenge(codeVerifier)
		state := generateState()

		// Store PKCE parameters in session/cookie (simplified for now)
		// In production, you'd store this in a secure session store
		http.SetCookie(w, &http.Cookie{
			Name:     "code_verifier",
			Value:    codeVerifier,
			Path:     "/",
			HttpOnly: true,
			Secure:   !a.config.DevMode,
			SameSite: http.SameSiteStrictMode,
			MaxAge:   600, // 10 minutes
		})

		http.SetCookie(w, &http.Cookie{
			Name:     "oauth_state",
			Value:    state,
			Path:     "/",
			HttpOnly: true,
			Secure:   !a.config.DevMode,
			SameSite: http.SameSiteStrictMode,
			MaxAge:   600, // 10 minutes
		})

		// Build authorization URL with PKCE
		authURL := a.oauth2.AuthCodeURL(state,
			oauth2.SetAuthURLParam("code_challenge", codeChallenge),
			oauth2.SetAuthURLParam("code_challenge_method", "S256"),
		)

		response := map[string]interface{}{
			"redirect_url": authURL,
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(response)
		return
	}

	http.Error(w, "Authentication provider not configured", http.StatusBadRequest)
}

// HandleOAuthCallback handles OAuth provider callback
func (a *AuthService) HandleOAuthCallback(w http.ResponseWriter, r *http.Request) {
	code := r.URL.Query().Get("code")
	state := r.URL.Query().Get("state")

	// Handle mock authentication in dev mode
	if a.config.DevMode && a.config.Provider == "mock" && code != "" {
		mockToken, err := a.CreateMockJWT("dev-user", "dev@localhost", "Development User")
		if err != nil {
			http.Error(w, "Failed to create mock token", http.StatusInternalServerError)
			return
		}

		response := map[string]interface{}{
			"token":    mockToken,
			"dev_mode": true,
			"message":  "Mock OAuth callback successful",
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(response)
		return
	}

	// Handle Authentik OIDC callback
	if a.config.Provider == "authentik" && a.oauth2 != nil {
		// Validate state parameter
		stateCookie, err := r.Cookie("oauth_state")
		if err != nil || stateCookie.Value != state {
			http.Error(w, "Invalid state parameter", http.StatusBadRequest)
			return
		}

		// Get code verifier from cookie
		verifierCookie, err := r.Cookie("code_verifier")
		if err != nil {
			http.Error(w, "Missing code verifier", http.StatusBadRequest)
			return
		}

		// Exchange code for token with PKCE
		ctx := context.Background()
		token, err := a.oauth2.Exchange(ctx, code,
			oauth2.SetAuthURLParam("code_verifier", verifierCookie.Value),
		)
		if err != nil {
			http.Error(w, fmt.Sprintf("Token exchange failed: %v", err), http.StatusBadRequest)
			return
		}

		// Extract ID token
		rawIDToken, ok := token.Extra("id_token").(string)
		if !ok {
			http.Error(w, "No ID token in response", http.StatusBadRequest)
			return
		}

		// Clear cookies
		http.SetCookie(w, &http.Cookie{Name: "oauth_state", MaxAge: -1, Path: "/"})
		http.SetCookie(w, &http.Cookie{Name: "code_verifier", MaxAge: -1, Path: "/"})

		response := map[string]interface{}{
			"token": rawIDToken,
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(response)
		return
	}

	http.Error(w, "Authentication provider not configured", http.StatusBadRequest)
}