Lamont/core/auth.py

import datetime
import os

import bcrypt
import jwt
from jwt.exceptions import ExpiredSignatureError, InvalidTokenError

import core.api_keys as apiKeys
import core.postgres as postgres
import core.users as users


DEFAULT_TOKEN_LIFETIME = datetime.timedelta(hours=1)


def decodeLoginToken(loginToken):
    secret = os.getenv("JWT_SECRET")
    if not secret or not isinstance(loginToken, str):
        return None
    try:
        payload = jwt.decode(loginToken, secret, algorithms=["HS256"])
    except (ExpiredSignatureError, InvalidTokenError):
        return None

    userUUID = payload.get("sub")
    if not isinstance(userUUID, str) or not users.doesUserUUIDExist(userUUID):
        return None
    return payload


def createLoginToken(userUUID, name=None, expiresIn=None, extraClaims=None):
    secret = os.getenv("JWT_SECRET")
    if not secret:
        raise RuntimeError("JWT_SECRET is required")
    if not userUUID or not users.doesUserUUIDExist(userUUID):
        raise ValueError("user does not exist")

    now = datetime.datetime.now(datetime.timezone.utc)
    if expiresIn is None:
        expiresIn = DEFAULT_TOKEN_LIFETIME
    elif isinstance(expiresIn, (int, float)):
        expiresIn = datetime.timedelta(seconds=expiresIn)
    if not isinstance(expiresIn, datetime.timedelta) or expiresIn.total_seconds() <= 0:
        raise ValueError("token lifetime must be positive")

    payload = {
        "sub": str(userUUID),
        "name": name if name is not None else users.getUserFirstName(userUUID),
        "iat": now,
        "exp": now + expiresIn,
    }
    if extraClaims is not None:
        if not isinstance(extraClaims, dict):
            raise ValueError("extra claims must be an object")
        protectedClaims = {"sub", "iat", "exp"}
        payload.update(
            {
                key: value
                for key, value in extraClaims.items()
                if key not in protectedClaims
            }
        )
    return jwt.encode(payload, secret, algorithm="HS256")


def verifyLoginToken(login_token, username=False, userUUID=False):
    if username:
        userUUID = users.getUserUUID(username)

    if userUUID:
        decodedToken = decodeLoginToken(login_token)
        if decodedToken and decodedToken.get("sub") == str(userUUID):
            return True
    return False


def getUserpasswordHash(userUUID):
    user = postgres.select_one("users", {"id": userUUID})
    if user:
        pw_hash = user.get("password_hashed")
        if isinstance(pw_hash, memoryview):
            return bytes(pw_hash)
        return pw_hash
    return None


def getLoginToken(username, password):
    if not isinstance(username, str) or not isinstance(password, str):
        return False
    userUUID = users.getUserUUID(username)
    if userUUID:
        formattedPass = password.encode("utf-8")
        if not formattedPass or len(formattedPass) > users.MAX_PASSWORD_BYTES:
            return False
        usersHashedPassword = getUserpasswordHash(userUUID)
        if not usersHashedPassword:
            return False
        try:
            if bcrypt.checkpw(formattedPass, usersHashedPassword):
                return createLoginToken(userUUID)
        except (TypeError, ValueError):
            return False
    return False


def unregisterUser(userUUID, password):
    pw_hash = getUserpasswordHash(userUUID)
    if not pw_hash or not isinstance(password, str):
        return False
    formattedPassword = password.encode("utf-8")
    if not formattedPassword or len(formattedPassword) > users.MAX_PASSWORD_BYTES:
        return False
    try:
        if bcrypt.checkpw(formattedPassword, pw_hash):
            return users.deleteUser(userUUID)
    except (TypeError, ValueError):
        return False
    return False


def getBearerToken(authorization):
    if not isinstance(authorization, str):
        return None
    scheme, separator, token = authorization.strip().partition(" ")
    if not separator or scheme.lower() != "bearer" or not token.strip():
        return None
    return token.strip()


def authenticateBearerToken(
    authorization,
    requiredScopes=None,
    allowUser=True,
    allowService=True,
):
    token = getBearerToken(authorization)
    if not token:
        return None

    if token.count(".") == 2:
        payload = decodeLoginToken(token)
        if payload:
            if not allowUser or requiredScopes:
                return None
            return {
                "type": "user",
                "authentication": "jwt",
                "user_uuid": payload["sub"],
                "service_name": None,
                "api_key_id": None,
                "scopes": [],
                "can_manage_api_keys": True,
                "claims": payload,
            }

    principal = apiKeys.authenticateApiKey(token, requiredScopes=requiredScopes)
    if not principal:
        return None
    if principal["type"] == apiKeys.USER_KEY_TYPE and not allowUser:
        return None
    if principal["type"] == apiKeys.SERVICE_KEY_TYPE and not allowService:
        return None
    return principal


def isUserPrincipal(principal, userUUID=None, requireLogin=False):
    if not principal or principal.get("type") != "user":
        return False
    if userUUID is not None and str(principal.get("user_uuid")) != str(userUUID):
        return False
    if requireLogin and principal.get("authentication") != "jwt":
        return False
    return True


def hasServiceScope(principal, scope):
    if not principal or principal.get("type") != "service":
        return False
    scopes = principal.get("scopes") or []
    return "*" in scopes or scope in scopes