python-kmsauth/kmsauth/services.py at master · lyft/python-kmsauth · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
"""Module for accessing boto3 clients, resources and sessions."""

import boto3
import botocore
import logging

CLIENT_CACHE = {}
RESOURCE_CACHE = {}


def get_boto_client(
        client,
        region=None,
        aws_access_key_id=None,
        aws_secret_access_key=None,
        aws_session_token=None,
        endpoint_url=None,
        max_pool_connections=None,
        connect_timeout=None,
        read_timeout=None,
        ):
    """Get a boto3 client connection."""
    cache_key = '{0}:{1}:{2}:{3}'.format(
        client,
        region,
        aws_access_key_id,
        endpoint_url or ''
    )
    if not aws_session_token:
        if cache_key in CLIENT_CACHE:
            return CLIENT_CACHE[cache_key]
    session = get_boto_session(
        region,
        aws_access_key_id,
        aws_secret_access_key,
        aws_session_token
    )
    if not session:
        logging.error("Failed to get {0} client.".format(client))
        return None

    # do not explicitly set any params as None
    config_params = dict(
        max_pool_connections=max_pool_connections,
        connect_timeout=connect_timeout,
        read_timeout=read_timeout,
    )
    config_params = {k: v for (k, v) in config_params.items() if v is not None}
    CLIENT_CACHE[cache_key] = session.client(
        client,
        endpoint_url=endpoint_url,
        config=botocore.config.Config(**config_params)
    )
    return CLIENT_CACHE[cache_key]


def get_boto_resource(
        resource,
        region=None,
        aws_access_key_id=None,
        aws_secret_access_key=None,
        aws_session_token=None,
        endpoint_url=None
        ):
    """Get a boto resource connection."""
    cache_key = '{0}:{1}:{2}:{3}'.format(
        resource,
        region,
        aws_access_key_id,
        endpoint_url or ''
    )
    if not aws_session_token:
        if cache_key in RESOURCE_CACHE:
            return RESOURCE_CACHE[cache_key]
    session = get_boto_session(
        region,
        aws_access_key_id,
        aws_secret_access_key,
        aws_session_token
    )
    if not session:
        logging.error("Failed to get {0} resource.".format(resource))
        return None

    RESOURCE_CACHE[cache_key] = session.resource(
        resource,
        endpoint_url=endpoint_url
    )
    return RESOURCE_CACHE[cache_key]


def get_boto_session(
        region,
        aws_access_key_id=None,
        aws_secret_access_key=None,
        aws_session_token=None
        ):
    """Get a boto3 session."""
    return boto3.session.Session(
        region_name=region,
        aws_secret_access_key=aws_secret_access_key,
        aws_access_key_id=aws_access_key_id,
        aws_session_token=aws_session_token
    )