mirror of https://github.com/kortix-ai/suna.git
592 lines
20 KiB
Python
592 lines
20 KiB
Python
import sentry
|
|
from fastapi import HTTPException, Request, Header
|
|
from typing import Optional
|
|
import jwt
|
|
from jwt.exceptions import PyJWTError
|
|
from utils.logger import structlog
|
|
from utils.config import config
|
|
import os
|
|
import base64
|
|
import hashlib
|
|
import hmac
|
|
from services.supabase import DBConnection
|
|
from services import redis
|
|
|
|
async def verify_admin_api_key(x_admin_api_key: Optional[str] = Header(None)):
|
|
if not config.KORTIX_ADMIN_API_KEY:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail="Admin API key not configured on server"
|
|
)
|
|
|
|
if not x_admin_api_key:
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Admin API key required. Include X-Admin-Api-Key header."
|
|
)
|
|
|
|
if x_admin_api_key != config.KORTIX_ADMIN_API_KEY:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="Invalid admin API key"
|
|
)
|
|
|
|
return True
|
|
|
|
def _verify_jwt_signature(token: str, secret: str) -> dict:
|
|
"""
|
|
Securely verify JWT signature using the Supabase JWT secret.
|
|
|
|
Args:
|
|
token: The JWT token to verify
|
|
secret: The Supabase JWT secret for signature verification
|
|
|
|
Returns:
|
|
dict: The decoded JWT payload if signature is valid
|
|
|
|
Raises:
|
|
PyJWTError: If signature verification fails or token is invalid
|
|
"""
|
|
try:
|
|
# Decode and verify the JWT signature using HS256 algorithm (Supabase default)
|
|
payload = jwt.decode(
|
|
token,
|
|
secret,
|
|
algorithms=["HS256"],
|
|
options={
|
|
"verify_signature": True,
|
|
"verify_exp": True,
|
|
"verify_aud": False, # Supabase JWTs may not always have audience
|
|
"verify_iss": False # Supabase JWTs may not always have issuer
|
|
}
|
|
)
|
|
return payload
|
|
except PyJWTError as e:
|
|
structlog.get_logger().warning(f"JWT signature verification failed: {e}")
|
|
raise
|
|
|
|
def _decode_jwt_safely(token: str) -> dict:
|
|
"""
|
|
Safely decode and verify a JWT token with proper signature verification.
|
|
Falls back to no verification only if JWT secret is not configured.
|
|
|
|
Args:
|
|
token: The JWT token to decode
|
|
|
|
Returns:
|
|
dict: The decoded JWT payload
|
|
|
|
Raises:
|
|
PyJWTError: If token is invalid or signature verification fails
|
|
"""
|
|
jwt_secret = config.SUPABASE_JWT_SECRET
|
|
|
|
if jwt_secret:
|
|
# Production mode: Verify signature
|
|
structlog.get_logger().debug("Verifying JWT signature")
|
|
return _verify_jwt_signature(token, jwt_secret)
|
|
else:
|
|
# Development mode: Log warning and decode without verification
|
|
structlog.get_logger().warning(
|
|
"JWT_SECRET not configured - using insecure JWT decoding. "
|
|
"This should only be used in development!"
|
|
)
|
|
return jwt.decode(token, options={"verify_signature": False})
|
|
|
|
async def _get_user_id_from_account_cached(account_id: str) -> Optional[str]:
|
|
"""
|
|
Get user_id from account_id with Redis caching for performance
|
|
|
|
Args:
|
|
account_id: The account ID to look up
|
|
|
|
Returns:
|
|
str: The primary owner user ID, or None if not found
|
|
"""
|
|
cache_key = f"account_user:{account_id}"
|
|
|
|
try:
|
|
# Check Redis cache first
|
|
redis_client = await redis.get_client()
|
|
cached_user_id = await redis_client.get(cache_key)
|
|
if cached_user_id:
|
|
return cached_user_id.decode('utf-8') if isinstance(cached_user_id, bytes) else cached_user_id
|
|
except Exception as e:
|
|
structlog.get_logger().warning(f"Redis cache lookup failed for account {account_id}: {e}")
|
|
|
|
try:
|
|
# Fallback to database
|
|
db = DBConnection()
|
|
await db.initialize()
|
|
client = await db.client
|
|
|
|
user_result = await client.schema('basejump').table('accounts').select(
|
|
'primary_owner_user_id'
|
|
).eq('id', account_id).limit(1).execute()
|
|
|
|
if user_result.data:
|
|
user_id = user_result.data[0]['primary_owner_user_id']
|
|
|
|
# Cache the result for 5 minutes
|
|
try:
|
|
await redis_client.setex(cache_key, 300, user_id)
|
|
except Exception as e:
|
|
structlog.get_logger().warning(f"Failed to cache user lookup: {e}")
|
|
|
|
return user_id
|
|
|
|
return None
|
|
|
|
except Exception as e:
|
|
structlog.get_logger().error(f"Database lookup failed for account {account_id}: {e}")
|
|
return None
|
|
|
|
# This function extracts the user ID from Supabase JWT
|
|
async def verify_and_get_user_id_from_jwt(request: Request) -> str:
|
|
"""
|
|
Extract and verify the user ID from the JWT in the Authorization header or API key.
|
|
|
|
This function is used as a dependency in FastAPI routes to ensure the user
|
|
is authenticated and to provide the user ID for authorization checks.
|
|
|
|
Supports authentication via:
|
|
1. X-API-Key header (public:secret key pairs from API keys table)
|
|
2. Authorization header with Bearer token (JWT from Supabase)
|
|
|
|
Args:
|
|
request: The FastAPI request object
|
|
|
|
Returns:
|
|
str: The user ID extracted from the JWT or API key
|
|
|
|
Raises:
|
|
HTTPException: If no valid token is found or if the token is invalid
|
|
"""
|
|
|
|
x_api_key = request.headers.get('x-api-key')
|
|
|
|
# Check for user API keys in the database
|
|
if x_api_key:
|
|
try:
|
|
# Parse the API key format: "pk_xxx:sk_xxx"
|
|
if ':' not in x_api_key:
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Invalid API key format. Expected format: pk_xxx:sk_xxx",
|
|
headers={"WWW-Authenticate": "Bearer"}
|
|
)
|
|
|
|
public_key, secret_key = x_api_key.split(':', 1)
|
|
|
|
from services.api_keys import APIKeyService
|
|
db = DBConnection()
|
|
await db.initialize()
|
|
api_key_service = APIKeyService(db)
|
|
|
|
validation_result = await api_key_service.validate_api_key(public_key, secret_key)
|
|
|
|
if validation_result.is_valid:
|
|
# Get user_id from account_id with caching
|
|
user_id = await _get_user_id_from_account_cached(str(validation_result.account_id))
|
|
|
|
if user_id:
|
|
sentry.sentry.set_user({ "id": user_id })
|
|
structlog.contextvars.bind_contextvars(
|
|
user_id=user_id,
|
|
auth_method="api_key",
|
|
api_key_id=str(validation_result.key_id),
|
|
public_key=public_key
|
|
)
|
|
return user_id
|
|
else:
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="API key account not found",
|
|
headers={"WWW-Authenticate": "Bearer"}
|
|
)
|
|
else:
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail=f"Invalid API key: {validation_result.error_message}",
|
|
headers={"WWW-Authenticate": "Bearer"}
|
|
)
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
structlog.get_logger().error(f"Error validating API key: {e}")
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="API key validation failed",
|
|
headers={"WWW-Authenticate": "Bearer"}
|
|
)
|
|
|
|
# Fall back to JWT authentication
|
|
auth_header = request.headers.get('Authorization')
|
|
|
|
if not auth_header or not auth_header.startswith('Bearer '):
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="No valid authentication credentials found",
|
|
headers={"WWW-Authenticate": "Bearer"}
|
|
)
|
|
|
|
token = auth_header.split(' ')[1]
|
|
|
|
try:
|
|
payload = _decode_jwt_safely(token)
|
|
user_id = payload.get('sub')
|
|
|
|
if not user_id:
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Invalid token payload",
|
|
headers={"WWW-Authenticate": "Bearer"}
|
|
)
|
|
|
|
sentry.sentry.set_user({ "id": user_id })
|
|
structlog.contextvars.bind_contextvars(
|
|
user_id=user_id,
|
|
auth_method="jwt"
|
|
)
|
|
return user_id
|
|
|
|
except PyJWTError:
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Invalid token",
|
|
headers={"WWW-Authenticate": "Bearer"}
|
|
)
|
|
|
|
|
|
async def get_user_id_from_stream_auth(
|
|
request: Request,
|
|
token: Optional[str] = None
|
|
) -> str:
|
|
"""
|
|
Extract and verify the user ID from multiple authentication methods.
|
|
This function is specifically designed for streaming endpoints that need to support both
|
|
header-based and query parameter-based authentication (for EventSource compatibility).
|
|
|
|
Supports authentication via:
|
|
1. X-API-Key header (public:secret key pairs from API keys table)
|
|
2. Authorization header with Bearer token (JWT from Supabase)
|
|
3. Query parameter token (JWT for EventSource compatibility)
|
|
|
|
Args:
|
|
request: The FastAPI request object
|
|
token: Optional JWT token from query parameters
|
|
|
|
Returns:
|
|
str: The user ID extracted from the authentication method
|
|
|
|
Raises:
|
|
HTTPException: If no valid token is found or if the token is invalid
|
|
"""
|
|
try:
|
|
# First, try the standard authentication (handles both API keys and Authorization header)
|
|
try:
|
|
return await verify_and_get_user_id_from_jwt(request)
|
|
except HTTPException:
|
|
# If standard auth fails, try query parameter JWT for EventSource compatibility
|
|
pass
|
|
|
|
# Try to get user_id from token in query param (for EventSource which can't set headers)
|
|
if token:
|
|
try:
|
|
# For Supabase JWT, verify signature and extract the user ID
|
|
payload = _decode_jwt_safely(token)
|
|
user_id = payload.get('sub')
|
|
if user_id:
|
|
sentry.sentry.set_user({ "id": user_id })
|
|
structlog.contextvars.bind_contextvars(
|
|
user_id=user_id,
|
|
auth_method="jwt_query"
|
|
)
|
|
return user_id
|
|
except Exception:
|
|
pass
|
|
|
|
# If we still don't have a user_id, return authentication error
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="No valid authentication credentials found",
|
|
headers={"WWW-Authenticate": "Bearer"}
|
|
)
|
|
except HTTPException:
|
|
# Re-raise HTTP exceptions as they are
|
|
raise
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
if "cannot schedule new futures after shutdown" in error_msg or "connection is closed" in error_msg:
|
|
raise HTTPException(
|
|
status_code=503,
|
|
detail="Server is shutting down"
|
|
)
|
|
else:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Error during authentication: {str(e)}"
|
|
)
|
|
|
|
async def get_optional_user_id(request: Request) -> Optional[str]:
|
|
"""
|
|
Extract the user ID from the JWT in the Authorization header if present,
|
|
but don't require authentication. Returns None if no valid token is found.
|
|
|
|
This function is used for endpoints that support both authenticated and
|
|
unauthenticated access (like public projects).
|
|
|
|
Args:
|
|
request: The FastAPI request object
|
|
|
|
Returns:
|
|
Optional[str]: The user ID extracted from the JWT, or None if no valid token
|
|
"""
|
|
auth_header = request.headers.get('Authorization')
|
|
|
|
if not auth_header or not auth_header.startswith('Bearer '):
|
|
return None
|
|
|
|
token = auth_header.split(' ')[1]
|
|
|
|
try:
|
|
# For Supabase JWT, verify signature and extract the user ID
|
|
payload = _decode_jwt_safely(token)
|
|
|
|
# Supabase stores the user ID in the 'sub' claim
|
|
user_id = payload.get('sub')
|
|
if user_id:
|
|
sentry.sentry.set_user({ "id": user_id })
|
|
structlog.contextvars.bind_contextvars(
|
|
user_id=user_id
|
|
)
|
|
|
|
return user_id
|
|
except PyJWTError:
|
|
return None
|
|
|
|
# Alias for consistency with other auth functions
|
|
get_optional_current_user_id_from_jwt = get_optional_user_id
|
|
|
|
async def verify_and_get_agent_authorization(client, agent_id: str, user_id: str) -> dict:
|
|
"""
|
|
Verify that a user has access to a specific agent based on ownership.
|
|
|
|
Args:
|
|
client: The Supabase client
|
|
agent_id: The agent ID to check access for
|
|
user_id: The user ID to check permissions for
|
|
|
|
Returns:
|
|
dict: Agent data if access is granted
|
|
|
|
Raises:
|
|
HTTPException: If the user doesn't have access to the agent or agent doesn't exist
|
|
"""
|
|
try:
|
|
agent_result = await client.table('agents').select('*').eq('agent_id', agent_id).eq('account_id', user_id).execute()
|
|
|
|
if not agent_result.data:
|
|
raise HTTPException(status_code=404, detail="Agent not found or access denied")
|
|
|
|
return agent_result.data[0]
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
structlog.error(f"Error verifying agent access for agent {agent_id}, user {user_id}: {str(e)}")
|
|
raise HTTPException(status_code=500, detail="Failed to verify agent access")
|
|
|
|
async def verify_and_authorize_thread_access(client, thread_id: str, user_id: str):
|
|
"""
|
|
Verify that a user has access to a specific thread based on account membership.
|
|
|
|
Args:
|
|
client: The Supabase client
|
|
thread_id: The thread ID to check access for
|
|
user_id: The user ID to check permissions for
|
|
|
|
Returns:
|
|
bool: True if the user has access
|
|
|
|
Raises:
|
|
HTTPException: If the user doesn't have access to the thread
|
|
"""
|
|
try:
|
|
# Query the thread to get account information
|
|
thread_result = await client.table('threads').select('*').eq('thread_id', thread_id).execute()
|
|
|
|
if not thread_result.data or len(thread_result.data) == 0:
|
|
raise HTTPException(status_code=404, detail="Thread not found")
|
|
|
|
thread_data = thread_result.data[0]
|
|
|
|
if thread_data['account_id'] == user_id:
|
|
return True
|
|
|
|
# Check if project is public
|
|
project_id = thread_data.get('project_id')
|
|
if project_id:
|
|
project_result = await client.table('projects').select('is_public').eq('project_id', project_id).execute()
|
|
if project_result.data and len(project_result.data) > 0:
|
|
if project_result.data[0].get('is_public'):
|
|
return True
|
|
|
|
account_id = thread_data.get('account_id')
|
|
# When using service role, we need to manually check account membership instead of using current_user_account_role
|
|
if account_id:
|
|
account_user_result = await client.schema('basejump').from_('account_user').select('account_role').eq('user_id', user_id).eq('account_id', account_id).execute()
|
|
if account_user_result.data and len(account_user_result.data) > 0:
|
|
return True
|
|
raise HTTPException(status_code=403, detail="Not authorized to access this thread")
|
|
except HTTPException:
|
|
# Re-raise HTTP exceptions as they are
|
|
raise
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
if "cannot schedule new futures after shutdown" in error_msg or "connection is closed" in error_msg:
|
|
raise HTTPException(
|
|
status_code=503,
|
|
detail="Server is shutting down"
|
|
)
|
|
else:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Error verifying thread access: {str(e)}"
|
|
)
|
|
|
|
# ============================================================================
|
|
# FastAPI Dependency Functions for Combined Authentication + Authorization
|
|
# ============================================================================
|
|
|
|
async def get_authorized_user_for_thread(
|
|
thread_id: str,
|
|
request: Request
|
|
) -> str:
|
|
"""
|
|
FastAPI dependency that verifies JWT and authorizes thread access.
|
|
|
|
Args:
|
|
thread_id: The thread ID to authorize access for
|
|
request: The FastAPI request object
|
|
|
|
Returns:
|
|
str: The authenticated and authorized user ID
|
|
|
|
Raises:
|
|
HTTPException: If authentication fails or user lacks thread access
|
|
"""
|
|
from services.supabase import DBConnection
|
|
|
|
# First, authenticate the user
|
|
user_id = await verify_and_get_user_id_from_jwt(request)
|
|
|
|
# Then, authorize thread access
|
|
db = DBConnection()
|
|
client = await db.client
|
|
await verify_and_authorize_thread_access(client, thread_id, user_id)
|
|
|
|
return user_id
|
|
|
|
async def get_authorized_user_for_agent(
|
|
agent_id: str,
|
|
request: Request
|
|
) -> tuple[str, dict]:
|
|
"""
|
|
FastAPI dependency that verifies JWT and authorizes agent access.
|
|
|
|
Args:
|
|
agent_id: The agent ID to authorize access for
|
|
request: The FastAPI request object
|
|
|
|
Returns:
|
|
tuple[str, dict]: The authenticated user ID and agent data
|
|
|
|
Raises:
|
|
HTTPException: If authentication fails or user lacks agent access
|
|
"""
|
|
from services.supabase import DBConnection
|
|
|
|
# First, authenticate the user
|
|
user_id = await verify_and_get_user_id_from_jwt(request)
|
|
|
|
# Then, authorize agent access and get agent data
|
|
db = DBConnection()
|
|
client = await db.client
|
|
agent_data = await verify_and_get_agent_authorization(client, agent_id, user_id)
|
|
|
|
return user_id, agent_data
|
|
|
|
class AuthorizedThreadAccess:
|
|
"""
|
|
FastAPI dependency that combines authentication and thread authorization.
|
|
|
|
Usage:
|
|
@router.get("/threads/{thread_id}/messages")
|
|
async def get_messages(
|
|
thread_id: str,
|
|
auth: AuthorizedThreadAccess = Depends()
|
|
):
|
|
user_id = auth.user_id # Authenticated and authorized user
|
|
"""
|
|
def __init__(self, user_id: str):
|
|
self.user_id = user_id
|
|
|
|
class AuthorizedAgentAccess:
|
|
"""
|
|
FastAPI dependency that combines authentication and agent authorization.
|
|
|
|
Usage:
|
|
@router.get("/agents/{agent_id}/config")
|
|
async def get_agent_config(
|
|
agent_id: str,
|
|
auth: AuthorizedAgentAccess = Depends()
|
|
):
|
|
user_id = auth.user_id # Authenticated and authorized user
|
|
agent_data = auth.agent_data # Agent data from authorization check
|
|
"""
|
|
def __init__(self, user_id: str, agent_data: dict):
|
|
self.user_id = user_id
|
|
self.agent_data = agent_data
|
|
|
|
async def require_thread_access(
|
|
thread_id: str,
|
|
request: Request
|
|
) -> AuthorizedThreadAccess:
|
|
"""
|
|
FastAPI dependency that verifies JWT and authorizes thread access.
|
|
|
|
Args:
|
|
thread_id: The thread ID from the path parameter
|
|
request: The FastAPI request object
|
|
|
|
Returns:
|
|
AuthorizedThreadAccess: Object containing authenticated user_id
|
|
|
|
Raises:
|
|
HTTPException: If authentication fails or user lacks thread access
|
|
"""
|
|
user_id = await get_authorized_user_for_thread(thread_id, request)
|
|
return AuthorizedThreadAccess(user_id)
|
|
|
|
async def require_agent_access(
|
|
agent_id: str,
|
|
request: Request
|
|
) -> AuthorizedAgentAccess:
|
|
"""
|
|
FastAPI dependency that verifies JWT and authorizes agent access.
|
|
|
|
Args:
|
|
agent_id: The agent ID from the path parameter
|
|
request: The FastAPI request object
|
|
|
|
Returns:
|
|
AuthorizedAgentAccess: Object containing user_id and agent_data
|
|
|
|
Raises:
|
|
HTTPException: If authentication fails or user lacks agent access
|
|
"""
|
|
user_id, agent_data = await get_authorized_user_for_agent(agent_id, request)
|
|
return AuthorizedAgentAccess(user_id, agent_data)
|
|
|