wroofauth/internal/machines/fesession/fe_session.go

package fesession

import (
	"context"
	"errors"
	"strings"

	"git.1in9.net/raider/wroofauth/internal/entities/user"
	"go.mongodb.org/mongo-driver/bson/primitive"
)

var (
	ErrIllegalStateAction = errors.New("illegal action for this state")
	ErrIllegalState       = errors.New("illegal state")
)

type SessionState string

const (
	SessionState_EMPTY                         SessionState = "EMPTY"
	SessionState_UNAUTHENTICATED               SessionState = "UNAUTHENTICATED"
	SessionState_AWAITING_FACTOR               SessionState = "AWAITING_FACTOR"
	SessionState_AUTHENTICATED_PENDING         SessionState = "AUTHENTICATED_PENDING"
	SessionState_AUTHENTICATED_FULLY           SessionState = "AUTHENTICATED_FULLY"
	SessionState_AUTHENTICATED_PASSWORD_CHANGE SessionState = "AUTHENTICATED_PASSWORD_CHANGE"
	SessionState_AUTHENTICATED_2FA_ENROLL      SessionState = "AUTHENTICATED_2FA_ENROLL"
	SessionState_AUTHENTICATED_REVIEW_TOS      SessionState = "AUTHENTICATED_REVIEW_TOS"
	SessionState_AUTHENTICATED_REVIEW_RECOVERY SessionState = "AUTHENTICATED_REVIEW_RECOVERY"
)

type AuthenticationMethod string

const (
	AuthenticationMethod_NONE     = "NONE"
	AuthenticationMethod_PASSWORD = "PASSWORD"
)

type SecondFactor string

const (
	SecondFactor_NONE = "NONE"
	SecondFactor_TOTP = "TOTP"
)

type FeSession struct {
	ID          primitive.ObjectID `bson:"_id"`
	LinkedToken primitive.ObjectID `bson:"token"`

	State SessionState `bson:"state"`

	AuthenticationMethod AuthenticationMethod `bson:"auth_method"`
	SecondFactor         SecondFactor         `bson:"second_factor"`

	User   *user.User          `bson:"-"`
	UserId *primitive.ObjectID `bson:"user"`
}

func NewSession(ctx context.Context, token primitive.ObjectID) *FeSession {
	return &FeSession{
		ID:                   primitive.NewObjectID(),
		State:                SessionState_EMPTY,
		AuthenticationMethod: AuthenticationMethod_NONE,
		SecondFactor:         SecondFactor_NONE,
	}
}

// s.Validate checks if the session is in a valid state
func (s *FeSession) Validate(ctx context.Context) error {
	if s.IsAnyAuthenticated() {
		if s.User == nil {
			// We can only be here if a user is set
			return ErrIllegalState
		}

		if s.AuthenticationMethod == AuthenticationMethod_NONE {
			// We can not be authenticated without any way we got here
			return ErrIllegalState
		}

		if s.SecondFactor == SecondFactor_NONE && s.User.Needs2FA() {
			// User has either just enabled 2FA, or something went horribly wrong.
			// Either way, throw away the session!
			return ErrIllegalState
		}
	}

	return nil
}

func (s *FeSession) IsAnyAuthenticated() bool {
	return strings.HasPrefix(string(s.State), "AUTHENTICATED_")
}

func (s *FeSession) performPreflight(ctx context.Context) error {
	// TODO: Do Preflight Checks.

	// TODO: Do PASSWORD_EXPIRED check

	// TODO: Do 2FA-Force-Enrolment check

	// TODO: Do TOS check

	// TODO: Do Reveiw Recovery check

	// TODO: Do Flag check

	// Preflight ok.
	s.State = SessionState_AUTHENTICATED_FULLY

	return nil
}

func (s *FeSession) Hydrate(ctx context.Context) error {
	if s.UserId != nil {
		var err error
		s.User, err = user.GetById(ctx, *s.UserId)
		if err != nil {
			return err
		}
	}

	return nil
}

func (s *FeSession) HandleIdentification(ctx context.Context, identification string) error {
	if s.State != SessionState_EMPTY {
		return ErrIllegalStateAction // This step may only run on EMPTY sessions
	}

	user, err := user.GetByIdentification(ctx, identification)
	if err != nil {
		return err
	}

	if user == nil {
		panic("no error, but nil user")
	}

	s.User = user
	s.UserId = &user.ID

	//TODO: Check for SAML

	s.State = SessionState_UNAUTHENTICATED

	return nil
}

func (s *FeSession) HandlePassword(ctx context.Context, password string) error {
	if s.State != SessionState_UNAUTHENTICATED {
		return ErrIllegalStateAction // This step may only run on UNAUTHENTICATED sessions
	}

	err := s.User.CheckPassword(password)
	if err != nil {
		return err
	}

	// Password Ok.
	s.AuthenticationMethod = AuthenticationMethod_PASSWORD

	if !s.User.Needs2FA() {
		// No 2fa, jump to AUTHENTICATED_PENDING for preflight
		s.State = SessionState_AUTHENTICATED_PENDING
		return s.performPreflight(ctx)
	}

	s.State = SessionState_AWAITING_FACTOR
	return nil
}

// TODO: Passkey action

func (s *FeSession) HandleTOTP(ctx context.Context, otp string) error {
	if s.State != SessionState_AWAITING_FACTOR {
		return ErrIllegalStateAction
	}

	err := s.User.ValidateTOTP(otp)
	if err != nil {
		return err
	}

	// OTP Ok.
	s.SecondFactor = SecondFactor_TOTP

	// Good to go for preflight.
	s.State = SessionState_AUTHENTICATED_PENDING
	return s.performPreflight(ctx)
}

func (s *FeSession) HandleLock(ctx context.Context) error {
	if !s.IsAnyAuthenticated() {
		return ErrIllegalStateAction
	}
	// TODO: Handle Lock
	return nil
}

func (s *FeSession) HandleLogout(ctx context.Context) error {
	if !s.IsAnyAuthenticated() {
		return ErrIllegalStateAction
	}
	// TODO: Handle Logout
	return nil
}

func (s *FeSession) Destroy(ctx context.Context) error {
	// TODO: Destroy Session
	return nil
}