gns3-server/gns3server/api/routes/controller/dependencies/authentication.py

#
# Copyright (C) 2020 GNS3 Technologies Inc.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import asyncio
import logging
import bcrypt

from fastapi import Request, Query, Depends, HTTPException, WebSocket, status
from fastapi.security import OAuth2PasswordBearer
from typing import Optional
from sqlalchemy import select

from gns3server import schemas
import gns3server.db.models as models
from gns3server.db.repositories.api_keys import ApiKeysRepository
from gns3server.db.repositories.users import UsersRepository
from gns3server.db.repositories.rbac import RbacRepository
from gns3server.services import auth_service
from .database import get_repository

log = logging.getLogger(__name__)
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/v3/access/users/login", auto_error=False)


async def get_user_from_token(
        bearer_token: str = Depends(oauth2_scheme),
        user_repo: UsersRepository = Depends(get_repository(UsersRepository)),
        api_keys_repo: ApiKeysRepository = Depends(get_repository(ApiKeysRepository)),
        token: Optional[str] = Query(None, include_in_schema=False)
) -> schemas.User:

    import time
    _t0 = time.time()
    log.info(f"[CTRL-TIMING] get_user_from_token ENTER bearer={bool(bearer_token)}")

    if bearer_token:
        # bearer token is used first, then any token passed as a URL parameter
        token = bearer_token

    if token is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Not authenticated",
            headers={"WWW-Authenticate": "Bearer"},
        )

    # API Key authentication
    if token.startswith("gns3_"):
        log.info(f"[CTRL-TIMING] get_user_from_token API_KEY auth elapsed={time.time()-_t0:.3f}s")
        query = select(models.ApiKey).where(models.ApiKey.revoked == False)
        result = await api_keys_repo._db_session.execute(query)
        api_keys_list = result.scalars().all()
        log.info(f"[CTRL-TIMING] get_user_from_token api_keys_count={len(api_keys_list)} elapsed={time.time()-_t0:.3f}s")
        for db_key in api_keys_list:
            # bcrypt.checkpw is CPU-bound and blocks the event loop; run in thread
            if await asyncio.to_thread(bcrypt.checkpw, token.encode(), db_key.key_hash.encode()):
                await api_keys_repo.update_last_used(db_key.api_key_id)
                user = await user_repo.get_user(db_key.user_id)
                if user:
                    if not user.is_active:
                        raise HTTPException(
                            status_code=status.HTTP_401_UNAUTHORIZED,
                            detail="Not an active user",
                            headers={"WWW-Authenticate": "Bearer"},
                        )
                    return user
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid API key",
            headers={"WWW-Authenticate": "Bearer"},
        )

    # JWT authentication
    token_data = auth_service.get_token_data(token)
    user = await user_repo.get_user_by_username(token_data.username)
    if user is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Could not validate credentials",
            headers={"WWW-Authenticate": "Bearer"},
        )
    if token_data.token_version != user.token_version:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=f"Token has been revoked for '{token_data.username}'",
            headers={"WWW-Authenticate": "Bearer"},
        )
    log.info(f"[CTRL-TIMING] get_user_from_token DONE elapsed={time.time()-_t0:.3f}s user={user.username}")
    return user


async def get_current_active_user(
        request: Request,
        current_user: schemas.User = Depends(get_user_from_token),
        rbac_repo: RbacRepository = Depends(get_repository(RbacRepository))
) -> schemas.User:

    # Super admin is always authorized
    if current_user.is_superadmin:
        return current_user

    if not current_user.is_active:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Not an active user",
            headers={"WWW-Authenticate": "Bearer"},
        )

    return current_user


async def get_current_active_user_from_websocket(
        websocket: WebSocket,
        token: str = Query(...),
        user_repo: UsersRepository = Depends(get_repository(UsersRepository)),
) -> Optional[schemas.User]:

    # Extract requested subprotocols from headers for proper WebSocket negotiation
    # This is critical for protocols like xpra that require specific subprotocols
    scope = websocket.scope
    headers = dict(scope.get("headers", []))
    requested_protocols_header = headers.get(b"sec-websocket-protocol", b"")
    requested_protocols = [p.decode().strip() for p in requested_protocols_header.split(b",") if p.strip()]

    # Accept the connection with the first requested subprotocol (if any)
    subprotocol = requested_protocols[0] if requested_protocols else None
    await websocket.accept(subprotocol=subprotocol)

    try:
        token_data = auth_service.get_token_data(token)
        user = await user_repo.get_user_by_username(token_data.username)

        if user is None:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=f"Could not validate credentials for '{token_data.username}'"
            )
        if token_data.token_version != user.token_version:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=f"Token has been revoked for '{token_data.username}'"
            )

        # Super admin is always authorized
        if user.is_superadmin:
            return user

        if not user.is_active:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=f"'{token_data.username}' is not an active user"
            )

        return user

    except HTTPException as e:
        err_msg = f"Could not authenticate while connecting to controller WebSocket: {e.detail}"
        websocket_error = {"action": "log.error", "event": {"message": err_msg}}
        await websocket.send_json(websocket_error)
        log.error(err_msg)
        return await websocket.close(code=1008)