breakpilot-core/consent-service/internal/handlers/oauth_handlers.go

package handlers

import (
	"context"
	"net/http"
	"strings"

	"github.com/breakpilot/consent-service/internal/middleware"
	"github.com/breakpilot/consent-service/internal/models"
	"github.com/breakpilot/consent-service/internal/services"
	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
)

// OAuthHandler handles OAuth 2.0 endpoints
type OAuthHandler struct {
	oauthService *services.OAuthService
	totpService  *services.TOTPService
	authService  *services.AuthService
}

// NewOAuthHandler creates a new OAuthHandler
func NewOAuthHandler(oauthService *services.OAuthService, totpService *services.TOTPService, authService *services.AuthService) *OAuthHandler {
	return &OAuthHandler{
		oauthService: oauthService,
		totpService:  totpService,
		authService:  authService,
	}
}

// ========================================
// OAuth 2.0 Authorization Code Flow
// ========================================

// Authorize handles the OAuth 2.0 authorization request
// GET /oauth/authorize
func (h *OAuthHandler) Authorize(c *gin.Context) {
	responseType := c.Query("response_type")
	clientID := c.Query("client_id")
	redirectURI := c.Query("redirect_uri")
	scope := c.Query("scope")
	state := c.Query("state")
	codeChallenge := c.Query("code_challenge")
	codeChallengeMethod := c.Query("code_challenge_method")

	// Validate response_type
	if responseType != "code" {
		c.JSON(http.StatusBadRequest, gin.H{
			"error":             "unsupported_response_type",
			"error_description": "Only 'code' response_type is supported",
		})
		return
	}

	// Validate client
	ctx := context.Background()
	client, err := h.oauthService.ValidateClient(ctx, clientID)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"error":             "invalid_client",
			"error_description": "Unknown or invalid client_id",
		})
		return
	}

	// Validate redirect_uri
	if err := h.oauthService.ValidateRedirectURI(client, redirectURI); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"error":             "invalid_request",
			"error_description": "Invalid redirect_uri",
		})
		return
	}

	// Validate scopes
	scopes, err := h.oauthService.ValidateScopes(client, scope)
	if err != nil {
		redirectWithError(c, redirectURI, "invalid_scope", "One or more requested scopes are invalid", state)
		return
	}

	// For public clients, PKCE is required
	if client.IsPublic && codeChallenge == "" {
		redirectWithError(c, redirectURI, "invalid_request", "PKCE code_challenge is required for public clients", state)
		return
	}

	// Get authenticated user
	userID, err := middleware.GetUserID(c)
	if err != nil || userID == uuid.Nil {
		// User not authenticated - redirect to login
		// Store authorization request in session and redirect to login
		c.JSON(http.StatusUnauthorized, gin.H{
			"error":             "login_required",
			"error_description": "User must be authenticated to authorize",
			"login_url":         "/auth/login",
		})
		return
	}

	// Generate authorization code
	code, err := h.oauthService.GenerateAuthorizationCode(
		ctx, client, userID, redirectURI, scopes, codeChallenge, codeChallengeMethod,
	)
	if err != nil {
		redirectWithError(c, redirectURI, "server_error", "Failed to generate authorization code", state)
		return
	}

	// Redirect with code
	redirectURL := redirectURI + "?code=" + code
	if state != "" {
		redirectURL += "&state=" + state
	}

	c.Redirect(http.StatusFound, redirectURL)
}

// Token handles the OAuth 2.0 token request
// POST /oauth/token
func (h *OAuthHandler) Token(c *gin.Context) {
	grantType := c.PostForm("grant_type")

	switch grantType {
	case "authorization_code":
		h.tokenAuthorizationCode(c)
	case "refresh_token":
		h.tokenRefreshToken(c)
	default:
		c.JSON(http.StatusBadRequest, gin.H{
			"error":             "unsupported_grant_type",
			"error_description": "Only 'authorization_code' and 'refresh_token' grant types are supported",
		})
	}
}

// tokenAuthorizationCode handles the authorization_code grant
func (h *OAuthHandler) tokenAuthorizationCode(c *gin.Context) {
	code := c.PostForm("code")
	clientID := c.PostForm("client_id")
	redirectURI := c.PostForm("redirect_uri")
	codeVerifier := c.PostForm("code_verifier")

	if code == "" || clientID == "" || redirectURI == "" {
		c.JSON(http.StatusBadRequest, gin.H{
			"error":             "invalid_request",
			"error_description": "Missing required parameters: code, client_id, redirect_uri",
		})
		return
	}

	// Validate client
	ctx := context.Background()
	client, err := h.oauthService.ValidateClient(ctx, clientID)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"error":             "invalid_client",
			"error_description": "Unknown or invalid client_id",
		})
		return
	}

	// For confidential clients, validate client_secret
	if !client.IsPublic {
		clientSecret := c.PostForm("client_secret")
		if err := h.oauthService.ValidateClientSecret(client, clientSecret); err != nil {
			c.JSON(http.StatusUnauthorized, gin.H{
				"error":             "invalid_client",
				"error_description": "Invalid client credentials",
			})
			return
		}
	}

	// Exchange authorization code for tokens
	tokenResponse, err := h.oauthService.ExchangeAuthorizationCode(ctx, code, clientID, redirectURI, codeVerifier)
	if err != nil {
		switch err {
		case services.ErrCodeExpired:
			c.JSON(http.StatusBadRequest, gin.H{
				"error":             "invalid_grant",
				"error_description": "Authorization code has expired",
			})
		case services.ErrCodeUsed:
			c.JSON(http.StatusBadRequest, gin.H{
				"error":             "invalid_grant",
				"error_description": "Authorization code has already been used",
			})
		case services.ErrPKCEVerifyFailed:
			c.JSON(http.StatusBadRequest, gin.H{
				"error":             "invalid_grant",
				"error_description": "PKCE verification failed",
			})
		default:
			c.JSON(http.StatusBadRequest, gin.H{
				"error":             "invalid_grant",
				"error_description": "Invalid authorization code",
			})
		}
		return
	}

	c.JSON(http.StatusOK, tokenResponse)
}

// tokenRefreshToken handles the refresh_token grant
func (h *OAuthHandler) tokenRefreshToken(c *gin.Context) {
	refreshToken := c.PostForm("refresh_token")
	clientID := c.PostForm("client_id")
	scope := c.PostForm("scope")

	if refreshToken == "" || clientID == "" {
		c.JSON(http.StatusBadRequest, gin.H{
			"error":             "invalid_request",
			"error_description": "Missing required parameters: refresh_token, client_id",
		})
		return
	}

	ctx := context.Background()

	// Refresh access token
	tokenResponse, err := h.oauthService.RefreshAccessToken(ctx, refreshToken, clientID, scope)
	if err != nil {
		switch err {
		case services.ErrInvalidScope:
			c.JSON(http.StatusBadRequest, gin.H{
				"error":             "invalid_scope",
				"error_description": "Requested scope exceeds original grant",
			})
		default:
			c.JSON(http.StatusBadRequest, gin.H{
				"error":             "invalid_grant",
				"error_description": "Invalid or expired refresh token",
			})
		}
		return
	}

	c.JSON(http.StatusOK, tokenResponse)
}

// Revoke handles token revocation
// POST /oauth/revoke
func (h *OAuthHandler) Revoke(c *gin.Context) {
	token := c.PostForm("token")
	tokenTypeHint := c.PostForm("token_type_hint")

	if token == "" {
		c.JSON(http.StatusBadRequest, gin.H{
			"error":             "invalid_request",
			"error_description": "Missing token parameter",
		})
		return
	}

	ctx := context.Background()
	_ = h.oauthService.RevokeToken(ctx, token, tokenTypeHint)

	// RFC 7009: Always return 200 OK
	c.Status(http.StatusOK)
}

// Introspect handles token introspection (for resource servers)
// POST /oauth/introspect
func (h *OAuthHandler) Introspect(c *gin.Context) {
	token := c.PostForm("token")

	if token == "" {
		c.JSON(http.StatusBadRequest, gin.H{
			"error":             "invalid_request",
			"error_description": "Missing token parameter",
		})
		return
	}

	ctx := context.Background()
	claims, err := h.oauthService.ValidateAccessToken(ctx, token)
	if err != nil {
		c.JSON(http.StatusOK, gin.H{"active": false})
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"active":    true,
		"sub":       (*claims)["sub"],
		"client_id": (*claims)["client_id"],
		"scope":     (*claims)["scope"],
		"exp":       (*claims)["exp"],
		"iat":       (*claims)["iat"],
		"iss":       (*claims)["iss"],
	})
}

// ========================================
// 2FA (TOTP) Endpoints
// ========================================

// Setup2FA initiates 2FA setup
// POST /auth/2fa/setup
func (h *OAuthHandler) Setup2FA(c *gin.Context) {
	userID, err := middleware.GetUserID(c)
	if err != nil || userID == uuid.Nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "Authentication required"})
		return
	}

	// Get user email
	ctx := context.Background()
	user, err := h.authService.GetUserByID(ctx, userID)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to get user"})
		return
	}

	// Setup 2FA
	response, err := h.totpService.Setup2FA(ctx, userID, user.Email)
	if err != nil {
		switch err {
		case services.ErrTOTPAlreadyEnabled:
			c.JSON(http.StatusConflict, gin.H{"error": "2FA is already enabled for this account"})
		default:
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to setup 2FA"})
		}
		return
	}

	c.JSON(http.StatusOK, response)
}

// Verify2FASetup verifies the 2FA setup with a code
// POST /auth/2fa/verify-setup
func (h *OAuthHandler) Verify2FASetup(c *gin.Context) {
	userID, err := middleware.GetUserID(c)
	if err != nil || userID == uuid.Nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "Authentication required"})
		return
	}

	var req models.Verify2FARequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid request body"})
		return
	}

	ctx := context.Background()
	err = h.totpService.Verify2FASetup(ctx, userID, req.Code)
	if err != nil {
		switch err {
		case services.ErrTOTPAlreadyEnabled:
			c.JSON(http.StatusConflict, gin.H{"error": "2FA is already enabled"})
		case services.ErrTOTPInvalidCode:
			c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid 2FA code"})
		default:
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to verify 2FA setup"})
		}
		return
	}

	c.JSON(http.StatusOK, gin.H{"message": "2FA enabled successfully"})
}

// Verify2FAChallenge verifies a 2FA challenge during login
// POST /auth/2fa/verify
func (h *OAuthHandler) Verify2FAChallenge(c *gin.Context) {
	var req models.Verify2FAChallengeRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid request body"})
		return
	}

	ctx := context.Background()
	var userID *uuid.UUID
	var err error

	if req.RecoveryCode != "" {
		// Verify with recovery code
		userID, err = h.totpService.VerifyChallengeWithRecoveryCode(ctx, req.ChallengeID, req.RecoveryCode)
	} else {
		// Verify with TOTP code
		userID, err = h.totpService.VerifyChallenge(ctx, req.ChallengeID, req.Code)
	}

	if err != nil {
		switch err {
		case services.ErrTOTPChallengeExpired:
			c.JSON(http.StatusGone, gin.H{"error": "2FA challenge has expired"})
		case services.ErrTOTPInvalidCode:
			c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid 2FA code"})
		case services.ErrRecoveryCodeInvalid:
			c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid recovery code"})
		default:
			c.JSON(http.StatusBadRequest, gin.H{"error": "2FA verification failed"})
		}
		return
	}

	// Get user and generate tokens
	user, err := h.authService.GetUserByID(ctx, *userID)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to get user"})
		return
	}

	// Generate access token
	accessToken, err := h.authService.GenerateAccessToken(user)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to generate token"})
		return
	}

	// Generate refresh token
	refreshToken, refreshTokenHash, err := h.authService.GenerateRefreshToken()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to generate refresh token"})
		return
	}

	// Store session
	ipAddress := middleware.GetClientIP(c)
	userAgent := middleware.GetUserAgent(c)

	// We need direct DB access for this, or we need to add a method to AuthService
	// For now, we'll return the tokens and let the caller handle session storage
	c.JSON(http.StatusOK, gin.H{
		"access_token":  accessToken,
		"refresh_token": refreshToken,
		"token_type":    "Bearer",
		"expires_in":    3600,
		"user": map[string]interface{}{
			"id":    user.ID,
			"email": user.Email,
			"name":  user.Name,
			"role":  user.Role,
		},
		"_session_hash": refreshTokenHash,
		"_ip":           ipAddress,
		"_user_agent":   userAgent,
	})
}

// Disable2FA disables 2FA for the current user
// POST /auth/2fa/disable
func (h *OAuthHandler) Disable2FA(c *gin.Context) {
	userID, err := middleware.GetUserID(c)
	if err != nil || userID == uuid.Nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "Authentication required"})
		return
	}

	var req models.Verify2FARequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid request body"})
		return
	}

	ctx := context.Background()
	err = h.totpService.Disable2FA(ctx, userID, req.Code)
	if err != nil {
		switch err {
		case services.ErrTOTPNotEnabled:
			c.JSON(http.StatusNotFound, gin.H{"error": "2FA is not enabled"})
		case services.ErrTOTPInvalidCode:
			c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid 2FA code"})
		default:
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to disable 2FA"})
		}
		return
	}

	c.JSON(http.StatusOK, gin.H{"message": "2FA disabled successfully"})
}

// Get2FAStatus returns the 2FA status for the current user
// GET /auth/2fa/status
func (h *OAuthHandler) Get2FAStatus(c *gin.Context) {
	userID, err := middleware.GetUserID(c)
	if err != nil || userID == uuid.Nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "Authentication required"})
		return
	}

	ctx := context.Background()
	status, err := h.totpService.GetStatus(ctx, userID)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to get 2FA status"})
		return
	}

	c.JSON(http.StatusOK, status)
}

// RegenerateRecoveryCodes generates new recovery codes
// POST /auth/2fa/recovery-codes
func (h *OAuthHandler) RegenerateRecoveryCodes(c *gin.Context) {
	userID, err := middleware.GetUserID(c)
	if err != nil || userID == uuid.Nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "Authentication required"})
		return
	}

	var req models.Verify2FARequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid request body"})
		return
	}

	ctx := context.Background()
	codes, err := h.totpService.RegenerateRecoveryCodes(ctx, userID, req.Code)
	if err != nil {
		switch err {
		case services.ErrTOTPNotEnabled:
			c.JSON(http.StatusNotFound, gin.H{"error": "2FA is not enabled"})
		case services.ErrTOTPInvalidCode:
			c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid 2FA code"})
		default:
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to regenerate recovery codes"})
		}
		return
	}

	c.JSON(http.StatusOK, gin.H{"recovery_codes": codes})
}

// ========================================
// Enhanced Login with 2FA
// ========================================

// LoginWith2FA handles login with optional 2FA
// POST /auth/login
func (h *OAuthHandler) LoginWith2FA(c *gin.Context) {
	var req models.LoginRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid request body"})
		return
	}

	ctx := context.Background()
	ipAddress := middleware.GetClientIP(c)
	userAgent := middleware.GetUserAgent(c)

	// Attempt login
	response, err := h.authService.Login(ctx, &req, ipAddress, userAgent)
	if err != nil {
		switch err {
		case services.ErrInvalidCredentials:
			c.JSON(http.StatusUnauthorized, gin.H{"error": "Invalid email or password"})
		case services.ErrAccountLocked:
			c.JSON(http.StatusForbidden, gin.H{"error": "Account is temporarily locked"})
		case services.ErrAccountSuspended:
			c.JSON(http.StatusForbidden, gin.H{"error": "Account is suspended"})
		default:
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Login failed"})
		}
		return
	}

	// Check if 2FA is enabled
	twoFactorEnabled, _ := h.totpService.IsTwoFactorEnabled(ctx, response.User.ID)

	if twoFactorEnabled {
		// Create 2FA challenge
		challengeID, err := h.totpService.CreateChallenge(ctx, response.User.ID, ipAddress, userAgent)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create 2FA challenge"})
			return
		}

		// Return 2FA required response
		c.JSON(http.StatusOK, gin.H{
			"requires_2fa":  true,
			"challenge_id":  challengeID,
			"message":       "2FA verification required",
		})
		return
	}

	// No 2FA required, return tokens
	c.JSON(http.StatusOK, gin.H{
		"requires_2fa":  false,
		"access_token":  response.AccessToken,
		"refresh_token": response.RefreshToken,
		"token_type":    "Bearer",
		"expires_in":    response.ExpiresIn,
		"user": map[string]interface{}{
			"id":    response.User.ID,
			"email": response.User.Email,
			"name":  response.User.Name,
			"role":  response.User.Role,
		},
	})
}

// ========================================
// Registration with mandatory 2FA setup
// ========================================

// RegisterWith2FA handles registration with mandatory 2FA setup
// POST /auth/register
func (h *OAuthHandler) RegisterWith2FA(c *gin.Context) {
	var req models.RegisterRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid request body"})
		return
	}

	ctx := context.Background()

	// Validate password strength
	if len(req.Password) < 8 {
		c.JSON(http.StatusBadRequest, gin.H{"error": "Password must be at least 8 characters"})
		return
	}

	// Register user
	user, verificationToken, err := h.authService.Register(ctx, &req)
	if err != nil {
		switch err {
		case services.ErrUserExists:
			c.JSON(http.StatusConflict, gin.H{"error": "A user with this email already exists"})
		default:
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Registration failed"})
		}
		return
	}

	// Setup 2FA immediately
	twoFAResponse, err := h.totpService.Setup2FA(ctx, user.ID, user.Email)
	if err != nil {
		// Non-fatal - user can set up 2FA later, but log it
		c.JSON(http.StatusCreated, gin.H{
			"message":            "Registration successful. Please verify your email.",
			"user_id":            user.ID,
			"verification_token": verificationToken, // In production, this would be sent via email
			"two_factor_setup":   nil,
			"two_factor_error":   "Failed to initialize 2FA. Please set it up in your account settings.",
		})
		return
	}

	c.JSON(http.StatusCreated, gin.H{
		"message":            "Registration successful. Please verify your email and complete 2FA setup.",
		"user_id":            user.ID,
		"verification_token": verificationToken, // In production, this would be sent via email
		"two_factor_setup": map[string]interface{}{
			"secret":          twoFAResponse.Secret,
			"qr_code":         twoFAResponse.QRCodeDataURL,
			"recovery_codes":  twoFAResponse.RecoveryCodes,
			"setup_required":  true,
			"setup_endpoint":  "/auth/2fa/verify-setup",
		},
	})
}

// ========================================
// OAuth Client Management (Admin)
// ========================================

// AdminCreateClient creates a new OAuth client
// POST /admin/oauth/clients
func (h *OAuthHandler) AdminCreateClient(c *gin.Context) {
	var req struct {
		Name         string   `json:"name" binding:"required"`
		Description  string   `json:"description"`
		RedirectURIs []string `json:"redirect_uris" binding:"required"`
		Scopes       []string `json:"scopes"`
		GrantTypes   []string `json:"grant_types"`
		IsPublic     bool     `json:"is_public"`
	}

	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid request body"})
		return
	}

	userID, _ := middleware.GetUserID(c)

	// Default scopes
	if len(req.Scopes) == 0 {
		req.Scopes = []string{"openid", "profile", "email"}
	}

	// Default grant types
	if len(req.GrantTypes) == 0 {
		req.GrantTypes = []string{"authorization_code", "refresh_token"}
	}

	ctx := context.Background()
	client, clientSecret, err := h.oauthService.CreateClient(
		ctx, req.Name, req.Description, req.RedirectURIs, req.Scopes, req.GrantTypes, req.IsPublic, &userID,
	)

	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create client"})
		return
	}

	response := gin.H{
		"client_id":     client.ClientID,
		"name":          client.Name,
		"redirect_uris": client.RedirectURIs,
		"scopes":        client.Scopes,
		"grant_types":   client.GrantTypes,
		"is_public":     client.IsPublic,
	}

	// Only show client_secret once for confidential clients
	if !client.IsPublic && clientSecret != "" {
		response["client_secret"] = clientSecret
		response["client_secret_warning"] = "Store this secret securely. It will not be shown again."
	}

	c.JSON(http.StatusCreated, response)
}

// AdminGetClients lists all OAuth clients
// GET /admin/oauth/clients
func (h *OAuthHandler) AdminGetClients(c *gin.Context) {
	// This would need a new method in OAuthService
	// For now, return a placeholder
	c.JSON(http.StatusOK, gin.H{
		"clients": []interface{}{},
		"message": "Client listing not yet implemented",
	})
}

// ========================================
// Helper Functions
// ========================================

func redirectWithError(c *gin.Context, redirectURI, errorCode, errorDescription, state string) {
	separator := "?"
	if strings.Contains(redirectURI, "?") {
		separator = "&"
	}

	redirectURL := redirectURI + separator + "error=" + errorCode + "&error_description=" + errorDescription
	if state != "" {
		redirectURL += "&state=" + state
	}

	c.Redirect(http.StatusFound, redirectURL)
}