secrets.py

import base64
import json
import shelve
import datetime
from typing import Dict

import boto3
from botocore.exceptions import BotoCoreError, ClientError


def get_secret_string(secret: str, region: str = 'us-east-1') -> str:
    now = datetime.datetime.now()

    with shelve.open('secrets_cache') as db:
        # Check if the secret is in the cache and less than 6 hours old
        if secret in db and now - db[secret]['time'] < datetime.timedelta(hours=6):
            return db[secret]['value']

        else:
            # Create a Secrets Manager client
            session = boto3.session.Session()
            client = session.client(
                service_name='secretsmanager',
                region_name=region  # specify your AWS region
            )

            try:
                get_secret_value_response = client.get_secret_value(SecretId=secret)
            except ClientError as e:
                # Handle exceptions related to the secret not existing, etc.
                raise e
            else:
                if 'SecretString' in get_secret_value_response:
                    secret_string = get_secret_value_response['SecretString']
                else:
                    # Binary secret data is base64-decoded and returned as a string
                    secret_string = base64.b64decode(get_secret_value_response['SecretBinary'])

                # Cache the secret and the current time
                db[secret] = {'value': secret_string, 'time': now}

                return secret_string


def get_secret_dict(secret: str, region: str = 'us-east-1') -> Dict[str, str]:
    secret_string = get_secret_string(secret)

    # Assuming that the secret string is a JSON string
    try:
        secret_dict = json.loads(secret_string)
    except json.JSONDecodeError as e:
        # Handle exceptions related to the secret string not being valid JSON
        raise e

    return secret_dict


def test_get_secret_string():
    secret = get_secret_string('dev/all/openai/cmwhee@gmail/api-key')
    assert len(secret) > 0
    assert type(secret) is str
    assert secret.startswith('sk-')


def test_get_secret_dict():
    secret = get_secret_dict('prod/services/twilio/APIKey')
    assert len(secret) > 0
    assert type(secret) is dict
    assert 'SID' in secret
    assert 'Secret' in secret
    assert secret['SID'].startswith('SK')
    assert secret['Secret'].startswith('HD')