imagini/internal/auth/auth.go

package auth

import (
	"errors"

    "github.com/google/uuid"
	log "github.com/sirupsen/logrus"
	"gorm.io/gorm"
	"reichard.io/imagini/internal/db"
	"reichard.io/imagini/internal/config"
	"reichard.io/imagini/internal/models"
	"reichard.io/imagini/internal/session"

	"encoding/json"
	"fmt"
	"time"

	"github.com/lestrrat-go/jwx/jwa"
	"github.com/lestrrat-go/jwx/jwt"
)

type AuthManager struct {
    DB      *db.DBManager
    Config  *config.Config
    Session *session.SessionManager
}

func NewMgr(db *db.DBManager, c *config.Config) *AuthManager {
    session := session.NewMgr()
    return &AuthManager{
        DB: db,
        Config: c,
        Session: session,
    }
}

func (auth *AuthManager) AuthenticateUser(creds models.APICredentials) (bool, models.User) {
    // By Username
    foundUser, err := auth.DB.User(models.User{Username: creds.User})
    if errors.Is(err, gorm.ErrRecordNotFound) {
        foundUser, err = auth.DB.User(models.User{Email: creds.User})
    }

    // Error Checking
    if errors.Is(err, gorm.ErrRecordNotFound) {
        log.Warn("[auth] User not found: ", creds.User)
        return false, foundUser
    } else if err != nil {
        log.Error(err)
        return false, foundUser
    }

    log.Info("[auth] Authenticating user: ", foundUser.Username)

    // Determine Type
    switch foundUser.AuthType {
    case "Local":
        return authenticateLocalUser(foundUser, creds.Password), foundUser
    case "LDAP":
        return authenticateLDAPUser(foundUser, creds.Password), foundUser
    default:
        return false, foundUser
    }
}


func (auth *AuthManager) getRole(user models.User) string {
    // TODO: Lookup role of user
    return "User"
}

func (auth *AuthManager) ValidateJWTAccessToken(accessJWT string) (jwt.Token, bool) {
    byteAccessJWT := []byte(accessJWT)
    verifiedToken, err := jwt.ParseBytes(byteAccessJWT, jwt.WithVerify(jwa.HS256, []byte(auth.Config.JWTSecret)))
    if err != nil {
        fmt.Println("failed to parse payload: ", err)
        return nil, false
    }
    return verifiedToken, true
}

func (auth *AuthManager) RevokeRefreshToken() {

}

func (auth *AuthManager) ValidateJWTRefreshToken(refreshJWT string) (jwt.Token, bool) {
    byteRefreshJWT := []byte(refreshJWT)

    // Acquire Relevant Device
    unverifiedToken, err := jwt.ParseBytes(byteRefreshJWT)
    did, ok := unverifiedToken.Get("did")
    if !ok {
        return nil, false
    }
    deviceID, err := uuid.Parse(fmt.Sprintf("%v", did))
    if err != nil {
        return nil, false
    }
    device, err := auth.DB.Device(models.Device{Base: models.Base{UUID: deviceID}})
    if err != nil {
        return nil, false
    }

    // Verify Token
    verifiedToken, err := jwt.ParseBytes(byteRefreshJWT, jwt.WithVerify(jwa.HS256, []byte(device.RefreshKey)))
    if err != nil {
        fmt.Println("failed to parse payload: ", err)
        return nil, false
    }
    return verifiedToken, true
}

func (auth *AuthManager) UpdateRefreshToken(deviceID string) error {
    // TODO:
    //  - Remove Refresh token from Session AND DB
    //  - Call CreateRefreshToken
    return nil
}

func (auth *AuthManager) CreateJWTRefreshToken(user models.User, device models.Device) (string, error) {
    // Acquire Refresh Key
    byteKey := []byte(device.RefreshKey)

    // Create New Token
    tm := time.Now()
    t := jwt.New()
    t.Set(`did`, device.UUID)                           // Device ID
    t.Set(jwt.SubjectKey, user.UUID)                    // User ID
    t.Set(jwt.AudienceKey, `imagini`)                   // App ID
    t.Set(jwt.IssuedAtKey, tm)                          // Issued At

    // TODO: Depends on Device
    t.Set(jwt.ExpirationKey, tm.Add(time.Hour * 24))    // 1 Day Access Key

    // Validate Token Creation
    _, err := json.MarshalIndent(t, "", "  ")
    if err != nil {
        fmt.Printf("failed to generate JSON: %s\n", err)
        return "", err
    }

    // Sign Token
    signed, err := jwt.Sign(t, jwa.HS256, byteKey)
    if err != nil {
        log.Printf("failed to sign token: %s", err)
        return "", err
    }

    // Return Token
    return string(signed), nil
}

func (auth *AuthManager) CreateJWTAccessToken(user models.User, device models.Device) (string, error) {
    // Acquire Role
    role := auth.getRole(user)

    // Create New Token
    tm := time.Now()
    t := jwt.New()
    t.Set(`did`, device.UUID)                           // Device ID
    t.Set(`role`, role)                                 // User Role (Admin / User)
    t.Set(jwt.SubjectKey, user.UUID)                    // User ID
    t.Set(jwt.AudienceKey, `imagini`)                   // App ID
    t.Set(jwt.IssuedAtKey, tm)                          // Issued At
    t.Set(jwt.ExpirationKey, tm.Add(time.Minute * 30))  // 30 Minute Access Key

    // Validate Token Creation
    _, err := json.MarshalIndent(t, "", "  ")
    if err != nil {
        fmt.Printf("failed to generate JSON: %s\n", err)
        return "", err
    }

    // Use Server Key
    byteKey := []byte(auth.Config.JWTSecret)

    // Sign Token
    signed, err := jwt.Sign(t, jwa.HS256, byteKey)
    if err != nil {
        log.Printf("failed to sign token: %s", err)
        return "", err
    }

    // Return Token
    return string(signed), nil
}