oc-auth/controllers/oauth2.go

package controllers

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"oc-auth/conf"
	"oc-auth/infrastructure"
	auth_connectors "oc-auth/infrastructure/auth_connector"
	"regexp"
	"strconv"
	"strings"
	"time"

	oclib "cloud.o-forge.io/core/oc-lib"
	"cloud.o-forge.io/core/oc-lib/models/peer"
	model "cloud.o-forge.io/core/oc-lib/models/peer"
	beego "github.com/beego/beego/v2/server/web"
)

// OAuthController handles OAuth2 login/consent provider endpoints
type OAuthController struct {
	beego.Controller
}

// @Title GetLogin
// @Description Hydra redirects here with a login_challenge. Returns challenge info or auto-accepts if session exists.
// @Param   login_challenge   query    string  true  "The login challenge from Hydra"
// @Success 200 {object} auth_connectors.LoginChallenge
// @Failure 400 missing login_challenge
// @Failure 500 internal error
// @router /login [get]
func (o *OAuthController) GetLogin() {
	logger := oclib.GetLogger()
	challenge := o.Ctx.Input.Query("login_challenge")
	if challenge == "" {
		o.Ctx.ResponseWriter.WriteHeader(400)
		o.Data["json"] = map[string]string{"error": "missing login_challenge parameter"}
		o.ServeJSON()
		return
	}

	if conf.GetConfig().Local {
		// In local mode, return a mock challenge for dev
		o.Data["json"] = &auth_connectors.LoginChallenge{
			Skip:      false,
			Challenge: challenge,
		}
		o.ServeJSON()
		return
	}

	loginChallenge, err := infrastructure.GetAuthConnector().GetLoginChallenge(challenge)
	if err != nil {
		logger.Error().Msg("Failed to get login challenge: " + err.Error())
		o.Ctx.ResponseWriter.WriteHeader(500)
		o.Data["json"] = map[string]string{"error": err.Error()}
		o.ServeJSON()
		return
	}

	// If skip is true, the user already has an active session — auto-accept
	if loginChallenge.Skip {
		redirect, err := infrastructure.GetAuthConnector().AcceptLogin(challenge, loginChallenge.Subject)
		if err != nil {
			logger.Error().Msg("Failed to auto-accept login: " + err.Error())
			o.Ctx.ResponseWriter.WriteHeader(500)
			o.Data["json"] = map[string]string{"error": err.Error()}
			o.ServeJSON()
			return
		}
		o.Data["json"] = redirect
		o.ServeJSON()
		return
	}

	// Return challenge info so frontend can render login form
	o.Data["json"] = loginChallenge
	o.ServeJSON()
}

// @Title PostLogin
// @Description Authenticate user via LDAP and accept Hydra login challenge
// @Param   body  body  auth_connectors.LoginRequest  true  "Login credentials and challenge"
// @Success 200 {object} auth_connectors.Redirect
// @Failure 401 invalid credentials
// @Failure 500 internal error
// @router /login [post]
func (o *OAuthController) Login() {
	logger := oclib.GetLogger()
	var req auth_connectors.LoginRequest
	if err := json.Unmarshal(o.Ctx.Input.CopyBody(10000000), &req); err != nil {
		o.Ctx.ResponseWriter.WriteHeader(400)
		o.Data["json"] = map[string]string{"error": "invalid request body"}
		o.ServeJSON()
		return
	}

	if req.Username == "" || req.Password == "" {
		o.Ctx.ResponseWriter.WriteHeader(400)
		o.Data["json"] = map[string]string{"error": "username and password are required"}
		o.ServeJSON()
		return
	}

	// Authenticate via LDAP
	ldap := auth_connectors.New()
	found, err := ldap.Authenticate(o.Ctx.Request.Context(), req.Username, req.Password)
	if err != nil || !found {
		logger.Error().Msg("LDAP authentication failed for user: " + req.Username)
		o.Ctx.ResponseWriter.WriteHeader(401)
		o.Data["json"] = map[string]string{"error": "invalid credentials"}
		o.ServeJSON()
		return
	}

	if conf.GetConfig().Local {
		// In local mode, return a mock token for dev
		t := oclib.NewRequest(oclib.LibDataEnum(oclib.PEER), "", "", []string{}, nil).Search(
			nil, fmt.Sprintf("%v", model.SELF.EnumIndex()), false)
		if t.Err == "" && len(t.Data) > 0 {
			p := t.Data[0].(*model.Peer)
			c := infrastructure.GetClaims().BuildConsentSession("local", req.Username, p)
			now := time.Now().UTC()
			now = now.Add(3600 * time.Second)
			c.Session.AccessToken["exp"] = now.Unix()
			b, _ := json.Marshal(c)
			token := &auth_connectors.Token{
				Active:      true,
				TokenType:   "Bearer",
				ExpiresIn:   3600,
				AccessToken: "localtoken." + base64.StdEncoding.EncodeToString(b),
			}
			o.Data["json"] = token
		} else {
			o.Ctx.ResponseWriter.WriteHeader(401)
			o.Data["json"] = map[string]string{"error": "peer not found"}
		}
		o.ServeJSON()
		return
	}

	if req.LoginChallenge == "" {
		o.Ctx.ResponseWriter.WriteHeader(400)
		o.Data["json"] = map[string]string{"error": "login_challenge is required in non-local mode"}
		o.ServeJSON()
		return
	}

	// Accept the login challenge with Hydra
	redirect, err := infrastructure.GetAuthConnector().AcceptLogin(req.LoginChallenge, req.Username)
	if err != nil {
		logger.Error().Msg("Failed to accept login: " + err.Error())
		o.Ctx.ResponseWriter.WriteHeader(500)
		o.Data["json"] = map[string]string{"error": err.Error()}
		o.ServeJSON()
		return
	}

	// Return redirect_to so the frontend follows the OAuth2 flow
	o.Data["json"] = redirect
	o.ServeJSON()
}

// @Title Consent
// @Description Hydra redirects here with a consent_challenge. Auto-accepts consent with user permissions.
// @Param   consent_challenge   query    string  true  "The consent challenge from Hydra"
// @Success 200 {object} auth_connectors.Redirect
// @Failure 400 missing consent_challenge
// @Failure 500 internal error
// @router /consent [get]
func (o *OAuthController) Consent() {
	logger := oclib.GetLogger()
	challenge := o.Ctx.Input.Query("consent_challenge")
	if challenge == "" {
		o.Ctx.ResponseWriter.WriteHeader(400)
		o.Data["json"] = map[string]string{"error": "missing consent_challenge parameter"}
		o.ServeJSON()
		return
	}

	// Get consent challenge details from Hydra
	consentChallenge, err := infrastructure.GetAuthConnector().GetConsentChallenge(challenge)
	if err != nil {
		logger.Error().Msg("Failed to get consent challenge: " + err.Error())
		o.Ctx.ResponseWriter.WriteHeader(500)
		o.Data["json"] = map[string]string{"error": err.Error()}
		o.ServeJSON()
		return
	}

	// Get self peer for signing
	pp := oclib.NewRequest(oclib.LibDataEnum(oclib.PEER), "", "", []string{}, nil).Search(
		nil, strconv.Itoa(peer.SELF.EnumIndex()), false)
	if len(pp.Data) == 0 || pp.Code >= 300 || pp.Err != "" {
		logger.Error().Msg("Self peer not found")
		o.Ctx.ResponseWriter.WriteHeader(500)
		o.Data["json"] = map[string]string{"error": "self peer not found"}
		o.ServeJSON()
		return
	}
	p := pp.Data[0].(*peer.Peer)

	// Extract client_id from consent challenge
	clientID := ""
	if consentChallenge.Client != nil {
		if cid, ok := consentChallenge.Client["client_id"].(string); ok {
			clientID = cid
		}
	}

	// Build consent session with user permissions and claims
	session := infrastructure.GetClaims().BuildConsentSession(clientID, consentChallenge.Subject, p)

	// Accept the consent challenge — grant all requested scopes
	redirect, err := infrastructure.GetAuthConnector().AcceptConsent(challenge, consentChallenge.RequestedScope, session)
	if err != nil {
		logger.Error().Msg("Failed to accept consent: " + err.Error())
		o.Ctx.ResponseWriter.WriteHeader(500)
		o.Data["json"] = map[string]string{"error": err.Error()}
		o.ServeJSON()
		return
	}

	// Return redirect_to (callback URL with authorization code)
	o.Data["json"] = redirect
	o.ServeJSON()
}

// @Title GetLogout
// @Description Hydra redirects here with a logout_challenge. Accepts the challenge and returns a redirect URL.
// @Param   logout_challenge   query    string  true  "The logout challenge from Hydra"
// @Success 200 {object} auth_connectors.Redirect
// @Failure 400 missing logout_challenge
// @Failure 500 internal error
// @router /logout [get]
func (o *OAuthController) GetLogout() {
	logger := oclib.GetLogger()
	challenge := o.Ctx.Input.Query("logout_challenge")
	if challenge == "" {
		o.Ctx.ResponseWriter.WriteHeader(400)
		o.Data["json"] = map[string]string{"error": "missing logout_challenge parameter"}
		o.ServeJSON()
		return
	}

	if conf.GetConfig().Local {
		o.Data["json"] = &auth_connectors.Redirect{RedirectTo: ""}
		o.ServeJSON()
		return
	}

	_, err := infrastructure.GetAuthConnector().GetLogoutChallenge(challenge)
	if err != nil {
		logger.Error().Msg("Failed to get logout challenge: " + err.Error())
		o.Ctx.ResponseWriter.WriteHeader(500)
		o.Data["json"] = map[string]string{"error": err.Error()}
		o.ServeJSON()
		return
	}

	redirect, err := infrastructure.GetAuthConnector().AcceptLogout(challenge)
	if err != nil {
		logger.Error().Msg("Failed to accept logout challenge: " + err.Error())
		o.Ctx.ResponseWriter.WriteHeader(500)
		o.Data["json"] = map[string]string{"error": err.Error()}
		o.ServeJSON()
		return
	}

	o.Data["json"] = redirect
	o.ServeJSON()
}

// @Title Logout
// @Description Revoke an OAuth2 token
// @Param   Authorization  header  string  false  "Bearer token"
// @Param   client_id      query   string  true   "The client_id"
// @Success 200 {object} auth_connectors.Token
// @router /logout [delete]
func (o *OAuthController) LogOut() {
	clientID := o.Ctx.Input.Query("client_id")
	reqToken := extractBearerToken(o.Ctx.Request)

	if conf.GetConfig().Local {
		o.Data["json"] = map[string]string{"status": "logged out"}
		o.ServeJSON()
		return
	}

	err := infrastructure.GetAuthConnector().RevokeToken(reqToken, clientID)
	if err != nil {
		o.Ctx.ResponseWriter.WriteHeader(500)
		o.Data["json"] = map[string]string{"error": err.Error()}
	} else {
		o.Data["json"] = &auth_connectors.Token{
			AccessToken: reqToken,
			Active:      false,
		}
	}
	o.ServeJSON()
}

// @Title Refresh
// @Description Exchange a refresh_token for a new token set
// @Param   body  body  object  true  "refresh_token and client_id"
// @Success 200 {object} auth_connectors.TokenResponse
// @Failure 401 invalid refresh token
// @router /refresh [post]
func (o *OAuthController) Refresh() {
	logger := oclib.GetLogger()
	var body struct {
		RefreshToken string `json:"refresh_token"`
		ClientID     string `json:"client_id"`
	}
	json.Unmarshal(o.Ctx.Input.CopyBody(100000), &body)

	if conf.GetConfig().Local {
		o.Data["json"] = map[string]string{"error": "refresh not supported in local mode"}
		o.Ctx.ResponseWriter.WriteHeader(400)
		o.ServeJSON()
		return
	}

	if body.RefreshToken == "" {
		o.Ctx.ResponseWriter.WriteHeader(400)
		o.Data["json"] = map[string]string{"error": "refresh_token is required"}
		o.ServeJSON()
		return
	}

	token, err := infrastructure.GetAuthConnector().RefreshToken(body.RefreshToken, body.ClientID)
	if err != nil {
		logger.Error().Msg("Failed to refresh token: " + err.Error())
		o.Ctx.ResponseWriter.WriteHeader(401)
		o.Data["json"] = map[string]string{"error": err.Error()}
	} else {
		o.Data["json"] = token
	}
	o.ServeJSON()
}

// @Title Introspect
// @Description Introspect a token — respects Hydra's response
// @Param   Authorization  header  string  false  "Bearer token"
// @Success 200 {object} auth_connectors.IntrospectResult
// @router /introspect [get]
func (o *OAuthController) Introspect() {
	reqToken := extractBearerToken(o.Ctx.Request)
	if reqToken == "" {
		o.Ctx.ResponseWriter.WriteHeader(401)
		o.Data["json"] = map[string]string{"error": "missing bearer token"}
		o.ServeJSON()
		return
	}

	if conf.GetConfig().Local {
		o.Data["json"] = &auth_connectors.IntrospectResult{Active: true}
		o.ServeJSON()
		return
	}

	result, err := infrastructure.GetAuthConnector().Introspect(reqToken)
	if err != nil {
		o.Ctx.ResponseWriter.WriteHeader(500)
		o.Data["json"] = map[string]string{"error": err.Error()}
	} else if !result.Active {
		o.Ctx.ResponseWriter.WriteHeader(401)
		o.Data["json"] = result
	} else {
		o.Data["json"] = result
	}
	o.ServeJSON()
}

var whitelist = []string{
	"/login",
	"/refresh",
	"/introspect",
	"/consent",
}

// @Title AuthForward
// @Description Forward auth for Traefik — validates JWT via Hydra introspection
// @Param   Authorization  header  string  false  "Bearer token"
// @Success 200 {string}
// @router /forward [get]
func (o *OAuthController) InternalAuthForward() {
	reqToken := o.Ctx.Request.Header.Get("Authorization")
	if reqToken == "" {
		for _, w := range whitelist {
			if strings.Contains(o.Ctx.Request.Header.Get("X-Forwarded-Uri"), w) {
				o.Ctx.ResponseWriter.WriteHeader(200)
				o.ServeJSON()
				return
			}
		}
		o.Ctx.ResponseWriter.WriteHeader(401)
		o.ServeJSON()
		return
	}
	splitToken := strings.Split(reqToken, "Bearer ")
	if len(splitToken) < 2 {
		reqToken = ""
	} else {
		reqToken = splitToken[1]
	}
	origin, publicKey, external := o.extractOrigin(o.Ctx.Request)
	if !infrastructure.GetAuthConnector().CheckAuthForward(
		reqToken, publicKey, origin,
		o.Ctx.Request.Header.Get("X-Forwarded-Method"),
		o.Ctx.Request.Header.Get("X-Forwarded-Uri"), external) && origin != "" && publicKey != "" {
		o.Ctx.ResponseWriter.WriteHeader(401)
		o.ServeJSON()
		return
	}
	o.ServeJSON()
}

func (o *OAuthController) extractOrigin(request *http.Request) (string, string, bool) {
	user, peerID, groups := oclib.ExtractTokenInfo(*request)
	external := true
	publicKey := ""
	origin := o.Ctx.Request.Header.Get("X-Forwarded-Host")
	if origin == "" {
		origin = o.Ctx.Request.Header.Get("Origin")
	}
	searchStr := origin
	r := regexp.MustCompile("(:[0-9]+)")
	t := r.FindString(searchStr)
	if t != "" {
		searchStr = strings.Replace(searchStr, t, "", -1)
	}
	pp := oclib.NewRequest(oclib.LibDataEnum(oclib.PEER), user, peerID, groups, nil).Search(nil, searchStr, false)
	if pp.Code != 200 || len(pp.Data) == 0 {
		return "", "", external
	}
	p := pp.Data[0].(*model.Peer)
	publicKey = p.PublicKey
	origin = p.APIUrl
	if origin != "" {
		if p.Relation == peer.SELF {
			external = false
		}
	} else {
		external = false
	}
	return origin, publicKey, external
}

// ExtractClient extracts the client_id from a JWT token.
// Supports both standard JWT (3 parts with base64 payload) and local dev tokens.
func ExtractClient(request http.Request) string {
	reqToken := request.Header.Get("Authorization")
	splitToken := strings.Split(reqToken, "Bearer ")
	if len(splitToken) < 2 {
		return ""
	}
	reqToken = splitToken[1]
	if reqToken == "" {
		return ""
	}

	// Try to decode as standard JWT (header.payload.signature)
	parts := strings.Split(reqToken, ".")
	if len(parts) >= 2 {
		// Decode the payload (second part of JWT)
		payload := parts[1]
		// Add padding if needed
		switch len(payload) % 4 {
		case 2:
			payload += "=="
		case 3:
			payload += "="
		}
		bytes, err := base64.URLEncoding.DecodeString(payload)
		if err != nil {
			// Try standard base64 for local dev tokens
			bytes, err = base64.StdEncoding.DecodeString(parts[len(parts)-1])
			if err != nil {
				return ""
			}
		}
		m := map[string]interface{}{}
		if err := json.Unmarshal(bytes, &m); err != nil {
			return ""
		}
		// Standard JWT: look for client_id in top-level or ext claims
		if cid, ok := m["client_id"].(string); ok {
			return cid
		}
		if ext, ok := m["ext"].(map[string]interface{}); ok {
			if cid, ok := ext["client_id"].(string); ok {
				return cid
			}
		}
		// Local dev token format: session.id_token.client_id
		if session, ok := m["session"].(map[string]interface{}); ok {
			if idToken, ok := session["id_token"].(map[string]interface{}); ok {
				if cid, ok := idToken["client_id"].(string); ok {
					return cid
				}
			}
		}
	}
	return ""
}

// extractBearerToken extracts the token from the Authorization header
func extractBearerToken(r *http.Request) string {
	reqToken := r.Header.Get("Authorization")
	splitToken := strings.Split(reqToken, "Bearer ")
	if len(splitToken) < 2 {
		return ""
	}
	return splitToken[1]
}