diff --git a/litellm/proxy/auth/rds_iam_token.py b/litellm/proxy/auth/rds_iam_token.py new file mode 100644 index 00000000000..f65fc4a99d7 --- /dev/null +++ b/litellm/proxy/auth/rds_iam_token.py @@ -0,0 +1,179 @@ +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 + + 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="") + return cleaned_token diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e1bc363d631..6bf81148fba 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -177,6 +177,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, @@ -228,6 +234,7 @@ def run_server( local, num_workers, test_async, + iam_token_db_auth, num_requests, use_queue, health, @@ -448,6 +455,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