From 1cc7c7fc59721076e505becf6380770336f98211 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 5 Aug 2024 17:35:48 -0700 Subject: [PATCH 1/3] feat(proxy_cli.py): support iam-based auth to rds Initial pr for iam-based auth support for rds --- litellm/proxy/auth/rds_iam_token.py | 11 +++++++++++ litellm/proxy/proxy_cli.py | 25 +++++++++++++++++++++++++ 2 files changed, 36 insertions(+) create mode 100644 litellm/proxy/auth/rds_iam_token.py diff --git a/litellm/proxy/auth/rds_iam_token.py b/litellm/proxy/auth/rds_iam_token.py new file mode 100644 index 00000000000..719e16ac054 --- /dev/null +++ b/litellm/proxy/auth/rds_iam_token.py @@ -0,0 +1,11 @@ +def generate_iam_auth_token(db_host, db_port, db_user) -> str: + from urllib.parse import quote + + import boto3 + + client = boto3.client("rds") + token = client.generate_db_auth_token( + DBHostname=db_host, Port=db_port, DBUsername=db_user + ) + cleaned_token = quote(token, safe="") + return cleaned_token diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 1d0eef6a0e8..af78085a4c1 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -171,6 +171,12 @@ def is_port_in_use(port): is_flag=True, help="Calls async endpoints /queue/requests and /queue/response", ) +@click.option( + "--iam_token_db_auth", + default=False, + is_flag=True, + help="Connects to RDS DB with IAM token", +) @click.option( "--num_requests", default=10, @@ -222,6 +228,7 @@ def run_server( local, num_workers, test_async, + iam_token_db_auth, num_requests, use_queue, health, @@ -442,6 +449,24 @@ def run_server( db_connection_pool_limit = 100 db_connection_timeout = 60 + ### GET DB TOKEN FOR IAM AUTH ### + + if iam_token_db_auth: + from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token + + db_host = os.getenv("DATABASE_HOST") + db_port = os.getenv("DATABASE_PORT") + db_user = os.getenv("DATABASE_USER") + db_name = os.getenv("DATABASE_NAME") + + token = generate_iam_auth_token( + db_host=db_host, db_port=db_port, db_user=db_user + ) + + # print(f"token: {token}") + _db_url = f"postgresql://{db_user}:{token}@{db_host}:{db_port}/{db_name}" + os.environ["DATABASE_URL"] = _db_url + ### DECRYPT ENV VAR ### from litellm.proxy.secret_managers.aws_secret_manager import decrypt_env_var From d29f7c0e067249a1ba988cfe17c26f7a76ced93f Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 5 Aug 2024 19:59:48 -0700 Subject: [PATCH 2/3] build(config.yml): pin prisma version --- .circleci/config.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 77942273be4..5712c71ca6b 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -48,7 +48,7 @@ jobs: pip install opentelemetry-sdk==1.25.0 pip install opentelemetry-exporter-otlp==1.25.0 pip install openai==1.34.0 - pip install prisma + pip install prisma==0.11.0 pip install "detect_secrets==1.5.0" pip install "httpx==0.24.1" pip install fastapi From f2fe887e73c8e73fae02f00188fb4792374518e2 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 5 Aug 2024 21:01:17 -0700 Subject: [PATCH 3/3] fix(rds_iam_token.py): support sts based auth --- litellm/proxy/auth/rds_iam_token.py | 172 +++++++++++++++++++++++++++- 1 file changed, 170 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/rds_iam_token.py b/litellm/proxy/auth/rds_iam_token.py index 719e16ac054..f65fc4a99d7 100644 --- a/litellm/proxy/auth/rds_iam_token.py +++ b/litellm/proxy/auth/rds_iam_token.py @@ -1,10 +1,178 @@ +import os +from typing import Optional, Union + +import httpx + + +def init_rds_client( + aws_access_key_id: Optional[str] = None, + aws_secret_access_key: Optional[str] = None, + aws_region_name: Optional[str] = None, + aws_session_name: Optional[str] = None, + aws_profile_name: Optional[str] = None, + aws_role_name: Optional[str] = None, + aws_web_identity_token: Optional[str] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, +): + from litellm import get_secret + + # check for custom AWS_REGION_NAME and use it if not passed to init_bedrock_client + litellm_aws_region_name = get_secret("AWS_REGION_NAME", None) + standard_aws_region_name = get_secret("AWS_REGION", None) + ## CHECK IS 'os.environ/' passed in + # Define the list of parameters to check + params_to_check = [ + aws_access_key_id, + aws_secret_access_key, + aws_region_name, + aws_session_name, + aws_profile_name, + aws_role_name, + aws_web_identity_token, + ] + + # Iterate over parameters and update if needed + for i, param in enumerate(params_to_check): + if param and param.startswith("os.environ/"): + params_to_check[i] = get_secret(param) + # Assign updated values back to parameters + ( + aws_access_key_id, + aws_secret_access_key, + aws_region_name, + aws_session_name, + aws_profile_name, + aws_role_name, + aws_web_identity_token, + ) = params_to_check + + ### SET REGION NAME + region_name = aws_region_name + if aws_region_name: + region_name = aws_region_name + elif litellm_aws_region_name: + region_name = litellm_aws_region_name + elif standard_aws_region_name: + region_name = standard_aws_region_name + else: + raise Exception( + "AWS region not set: set AWS_REGION_NAME or AWS_REGION env variable or in .env file", + ) + + import boto3 + + if isinstance(timeout, float): + config = boto3.session.Config(connect_timeout=timeout, read_timeout=timeout) + elif isinstance(timeout, httpx.Timeout): + config = boto3.session.Config( + connect_timeout=timeout.connect, read_timeout=timeout.read + ) + else: + config = boto3.session.Config() + + ### CHECK STS ### + if ( + aws_web_identity_token is not None + and aws_role_name is not None + and aws_session_name is not None + ): + try: + oidc_token = open(aws_web_identity_token).read() # check if filepath + except Exception: + oidc_token = get_secret(aws_web_identity_token) + + if oidc_token is None: + raise Exception( + "OIDC token could not be retrieved from secret manager.", + ) + + sts_client = boto3.client("sts") + + # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html + # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html + sts_response = sts_client.assume_role_with_web_identity( + RoleArn=aws_role_name, + RoleSessionName=aws_session_name, + WebIdentityToken=oidc_token, + DurationSeconds=3600, + ) + + client = boto3.client( + service_name="rds", + aws_access_key_id=sts_response["Credentials"]["AccessKeyId"], + aws_secret_access_key=sts_response["Credentials"]["SecretAccessKey"], + aws_session_token=sts_response["Credentials"]["SessionToken"], + region_name=region_name, + config=config, + ) + elif aws_role_name is not None and aws_session_name is not None: + # use sts if role name passed in + sts_client = boto3.client( + "sts", + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + ) + + sts_response = sts_client.assume_role( + RoleArn=aws_role_name, RoleSessionName=aws_session_name + ) + + client = boto3.client( + service_name="rds", + aws_access_key_id=sts_response["Credentials"]["AccessKeyId"], + aws_secret_access_key=sts_response["Credentials"]["SecretAccessKey"], + aws_session_token=sts_response["Credentials"]["SessionToken"], + region_name=region_name, + config=config, + ) + elif aws_access_key_id is not None: + # uses auth params passed to completion + # aws_access_key_id is not None, assume user is trying to auth using litellm.completion + + client = boto3.client( + service_name="rds", + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + region_name=region_name, + config=config, + ) + elif aws_profile_name is not None: + # uses auth values from AWS profile usually stored in ~/.aws/credentials + + client = boto3.Session(profile_name=aws_profile_name).client( + service_name="rds", + region_name=region_name, + config=config, + ) + else: + # aws_access_key_id is None, assume user is trying to auth using env variables + # boto3 automatically reads env variables + + client = boto3.client( + service_name="bedrock-runtime", + region_name=region_name, + config=config, + ) + + return client + + def generate_iam_auth_token(db_host, db_port, db_user) -> str: from urllib.parse import quote import boto3 - client = boto3.client("rds") - token = client.generate_db_auth_token( + boto_client = init_rds_client( + aws_region_name=os.getenv("AWS_REGION_NAME"), + aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"), + aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"), + aws_session_name=os.getenv("AWS_SESSION_NAME"), + aws_profile_name=os.getenv("AWS_PROFILE_NAME"), + aws_role_name=os.getenv("AWS_ROLE_NAME"), + aws_web_identity_token=os.getenv("AWS_WEB_IDENTITY_TOKEN"), + ) + + token = boto_client.generate_db_auth_token( DBHostname=db_host, Port=db_port, DBUsername=db_user ) cleaned_token = quote(token, safe="")