From 07d9cca2ab3dac1b41481b7c597088267afe0c6c Mon Sep 17 00:00:00 2001 From: Nate Barbettini Date: Wed, 28 Aug 2024 17:15:24 -0700 Subject: [PATCH] API key auth for Actors (#23) Protects Actors by requiring the Engine to send a token signed with the Arcade API key. Also, updates the FastAPI example so the Arcade API key is sent to the Engine in a chat request. --- arcade/arcade/actor/core/auth.py | 18 +++++---------- arcade/arcade/actor/fastapi/actor.py | 5 ++-- arcade/arcade/actor/fastapi/auth.py | 14 ++++------- arcade/arcade/cli/main.py | 23 ++++++++++++++++--- arcade/arcade/core/client.py | 4 +--- arcade/arcade/core/config.py | 4 ---- .../fastapi/arcade_example_fastapi/main.py | 12 +++++++--- 7 files changed, 44 insertions(+), 36 deletions(-) diff --git a/arcade/arcade/actor/core/auth.py b/arcade/arcade/actor/core/auth.py index 8216c6a4..e1494405 100644 --- a/arcade/arcade/actor/core/auth.py +++ b/arcade/arcade/actor/core/auth.py @@ -5,13 +5,12 @@ import jwt from arcade.core.config import config -TOKEN_VER = "1" # noqa: S105 Possible hardcoded password assigned (false positive) +SUPPORTED_TOKEN_VER = "1" # noqa: S105 Possible hardcoded password assigned (false positive) @dataclass class TokenValidationResult: valid: bool - api_key: str | None = None error: str | None = None @@ -19,25 +18,20 @@ class SigningAlgorithm(str, Enum): HS256 = "HS256" -def validate_token(token: str) -> TokenValidationResult: +def validate_engine_token(token: str) -> TokenValidationResult: try: payload = jwt.decode( token, - config.api.secret, + config.api.key, algorithms=[SigningAlgorithm.HS256], verify=True, - issuer=config.engine_url, audience="actor", ) except (jwt.ExpiredSignatureError, jwt.InvalidTokenError) as e: return TokenValidationResult(valid=False, error=str(e)) - api_key = payload.get("api_key") - if api_key != config.api.key: - return TokenValidationResult(valid=False, error="Invalid API key") - token_ver = payload.get("ver") - if token_ver != TOKEN_VER: - return TokenValidationResult(valid=False, error=f"Unknown token version: {token_ver}") + if token_ver != SUPPORTED_TOKEN_VER: + return TokenValidationResult(valid=False, error=f"Unsupported token version: {token_ver}") - return TokenValidationResult(valid=True, api_key=api_key) + return TokenValidationResult(valid=True) diff --git a/arcade/arcade/actor/fastapi/actor.py b/arcade/arcade/actor/fastapi/actor.py index 0fc30595..9831a039 100644 --- a/arcade/arcade/actor/fastapi/actor.py +++ b/arcade/arcade/actor/fastapi/actor.py @@ -1,13 +1,14 @@ import json from typing import Any, Callable -from fastapi import FastAPI, Request +from fastapi import Depends, FastAPI, Request from arcade.actor.core.base import ( BaseActor, Router, ) from arcade.actor.core.common import RequestData +from arcade.actor.fastapi.auth import validate_engine_request from arcade.actor.utils import is_async_callable @@ -39,7 +40,7 @@ class FastAPIRouter(Router): async def wrapped_handler( request: Request, - # api_key: str = Depends(get_api_key), # TODO re-enable when Engine supports auth + _: None = Depends(validate_engine_request), ) -> Any: body_str = await request.body() body_json = json.loads(body_str) if body_str else {} diff --git a/arcade/arcade/actor/fastapi/auth.py b/arcade/arcade/actor/fastapi/auth.py index c6ea41f2..babda1e2 100644 --- a/arcade/arcade/actor/fastapi/auth.py +++ b/arcade/arcade/actor/fastapi/auth.py @@ -1,19 +1,17 @@ -from typing import cast - from fastapi import Depends, HTTPException from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer -from arcade.actor.core.auth import validate_token +from arcade.actor.core.auth import validate_engine_token security = HTTPBearer() # Authorization: Bearer -# Dependency function to validate JWT and extract API key -async def get_api_key( +# Dependency function to validate JWT +async def validate_engine_request( credentials: HTTPAuthorizationCredentials = Depends(security), -) -> str: +) -> None: jwt: str = credentials.credentials - validation_result = validate_token(jwt) + validation_result = validate_engine_token(jwt) if not validation_result.valid: raise HTTPException( @@ -21,5 +19,3 @@ async def get_api_key( detail=f"Invalid token. Error: {validation_result.error}", headers={"WWW-Authenticate": "Bearer"}, ) - - return cast(str, validation_result.api_key) diff --git a/arcade/arcade/cli/main.py b/arcade/arcade/cli/main.py index ca9824df..54269b94 100644 --- a/arcade/arcade/cli/main.py +++ b/arcade/arcade/cli/main.py @@ -126,8 +126,18 @@ def run( tools = [catalog[tool]] if tool else list(catalog) - # allow user to specify the engine url? - client = EngineClient() + config = Config.load_from_file() + if not config.engine or not config.engine_url: + console.print("❌ Engine configuration not found or URL is missing.", style="bold red") + typer.Exit(code=1) + + if not config.api or not config.api.key: + console.print( + "❌ API configuration not found or key is missing. Please run `arcade login`.", + style="bold red", + ) + typer.Exit(code=1) + client = EngineClient(api_key=config.api.key, base_url=config.engine_url) # TODO better way of doing this tool_choice = "auto" if choice in ["execute", "generate"] else choice @@ -202,7 +212,14 @@ def chat( console.print("❌ Engine configuration not found or URL is missing.", style="bold red") typer.Exit(code=1) - client = EngineClient(base_url=config.engine_url) + if not config.api or not config.api.key: + console.print( + "❌ API configuration not found or key is missing. Please run `arcade login`.", + style="bold red", + ) + typer.Exit(code=1) + + client = EngineClient(api_key=config.api.key, base_url=config.engine_url) if config.user and config.user.email: user_email = config.user.email diff --git a/arcade/arcade/core/client.py b/arcade/arcade/core/client.py index 9c30e3ed..6cd04576 100644 --- a/arcade/arcade/core/client.py +++ b/arcade/arcade/core/client.py @@ -1,5 +1,4 @@ import json -import os from enum import Enum from typing import Any, Optional @@ -117,8 +116,7 @@ def get_tool_args(chat_completion: ChatCompletion) -> list[tuple[str, dict[str, class EngineClient: - def __init__(self, api_key: str | None = None, base_url: str | None = None): - api_key = os.environ["OPENAI_API_KEY"] if api_key is None else api_key + def __init__(self, api_key: str, base_url: str | None = None): self.client = OpenAI(api_key=api_key, base_url=base_url) def __getattr__(self, name: str) -> Any: diff --git a/arcade/arcade/core/config.py b/arcade/arcade/core/config.py index d2ee4712..885078ea 100644 --- a/arcade/arcade/core/config.py +++ b/arcade/arcade/core/config.py @@ -15,10 +15,6 @@ class ApiConfig(BaseModel): """ Arcade API key. """ - secret: str - """ - Arcade API secret. - """ class UserConfig(BaseModel): diff --git a/examples/fastapi/arcade_example_fastapi/main.py b/examples/fastapi/arcade_example_fastapi/main.py index 610473c9..d1c6d746 100644 --- a/examples/fastapi/arcade_example_fastapi/main.py +++ b/examples/fastapi/arcade_example_fastapi/main.py @@ -5,10 +5,17 @@ from pydantic import BaseModel from arcade_gmail.tools import gmail from arcade_github.tools import repo, user from arcade_slack.tools import chat +from arcade.core.config import config from arcade.actor.fastapi.actor import FastAPIActor -client = AsyncOpenAI(base_url="http://localhost:9099/v1") +if not config.api or not config.api.key: + raise ValueError("Arcade API key not set. Please run `arcade login`.") + +client = AsyncOpenAI( + api_key=config.api.key, + base_url="http://localhost:9099/v1", +) app = FastAPI() @@ -40,7 +47,6 @@ async def postChat(request: ChatRequest, tool_choice: str = "execute"): ], model="gpt-4o-mini", max_tokens=500, - # TODO tests for tool choice tools=[ "GetEmails", "WriteDraft", @@ -51,7 +57,7 @@ async def postChat(request: ChatRequest, tool_choice: str = "execute"): "SendMessageToChannel", ], tool_choice=tool_choice, - user="sam", + user=config.user.email if config.user else None, ) return raw_response.choices except Exception as e: