A CredentialsProvider class has been added to allow the user to add his own provider for password rotation

Author: barshaul

docs/examples/connection_examples.ipynb

@@ -97,6 +97,92 @@
     "user_connection.ping()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "## Connecting to a redis instance with AWS Secrets Manager credentials provider."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import redis\n",
+    "import boto3\n",
+    "import json\n",
+    "import cachetools.func\n",
+    "\n",
+    "sm_client = boto3.client('secretsmanager')\n",
+    "  \n",
+    "def sm_auth_provider(secret_id, version_id=None, version_stage='AWSCURRENT'):\n",
+    "    @cachetools.func.ttl_cache(maxsize=128, ttl=24 * 60 * 60) #24h\n",
+    "    def get_sm_user_credentials(secret_id, version_id, version_stage):\n",
+    "        secret = sm_client.get_secret_value(secret_id, version_id)\n",
+    "        return json.loads(secret['SecretString'])\n",
+    "    creds = get_sm_user_credentials(secret_id, version_id, version_stage)\n",
+    "    return creds['username'], creds['password']\n",
+    "\n",
+    "secret_id = \"EXAMPLE1-90ab-cdef-fedc-ba987SECRET1\"\n",
+    "creds_provider = redis.CredentialsProvider(supplier=sm_auth_provider, secret_id=secret_id)\n",
+    "user_connection = redis.Redis(host=\"localhost\", port=6379, credentials_provider=creds_provider)\n",
+    "user_connection.ping()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Connecting to a redis instance with ElastiCache IAM credentials provider."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import redis\n",
+    "import boto3\n",
+    "import cachetools.func\n",
+    "\n",
+    "ec_client = boto3.client('elasticache')\n",
+    "\n",
+    "def iam_auth_provider(user, endpoint, port=6379, region=\"us-east-1\"):\n",
+    "    @cachetools.func.ttl_cache(maxsize=128, ttl=15 * 60) # 15m\n",
+    "    def get_iam_auth_token(user, endpoint, port, region):\n",
+    "        return ec_client.generate_iam_auth_token(user, endpoint, port, region)\n",
+    "    iam_auth_token = get_iam_auth_token(endpoint, port, user, region)\n",
+    "    return iam_auth_token\n",
+    "\n",
+    "username = \"barshaul\"\n",
+    "endpoint = \"test-001.use1.cache.amazonaws.com\"\n",
+    "creds_provider = redis.CredentialsProvider(supplier=iam_auth_provider, user=username,\n",
+    "                                           endpoint=endpoint)\n",
+    "user_connection = redis.Redis(host=endpoint, port=6379, credentials_provider=creds_provider)\n",
+    "user_connection.ping()"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -176,4 +262,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}

redis/__init__.py

@@ -6,6 +6,7 @@
     BlockingConnectionPool,
     Connection,
     ConnectionPool,
+    CredentialsProvider,
     SSLConnection,
     UnixDomainSocketConnection,
 )
@@ -62,6 +63,7 @@ def int_or_str(value):
     "Connection",
     "ConnectionError",
     "ConnectionPool",
+    "CredentialsProvider",
     "DataError",
     "from_url",
     "InvalidResponse",

redis/client.py

@@ -936,6 +936,7 @@ def __init__(
         username=None,
         retry=None,
         redis_connect_func=None,
+        credentials_provider=None,
     ):
         """
         Initialize a new Redis client.
@@ -977,6 +978,7 @@ def __init__(
                 "health_check_interval": health_check_interval,
                 "client_name": client_name,
                 "redis_connect_func": redis_connect_func,
+                "credentials_provider": credentials_provider,
             }
             # based on input, setup appropriate connection args
             if unix_socket_path is not None:

redis/cluster.py

@@ -125,6 +125,7 @@ def parse_cluster_shards(resp, **options):
     "connection_class",
     "connection_pool",
     "client_name",
+    "credentials_provider",
     "db",
     "decode_responses",
     "encoding",

redis/connection.py

@@ -500,6 +500,42 @@ def read_response(self, disable_decoding=False):
     DefaultParser = PythonParser
 
 
+class CredentialsProvider:
+    def __init__(self, username="", password="", supplier=None, *args, **kwargs):
+        """
+        Initialize a new Credentials Provider.
+        :param supplier: a supplier function that returns the username and password.
+                         def supplier(arg1, arg2, ...) -> (username, password)
+                         For examples see examples/connection_examples.ipynb
+        :param args: arguments to pass to the supplier function
+        :param kwargs: keyword arguments to pass to the supplier function
+        """
+        self.username = username
+        self.password = password
+        self.supplier = supplier
+        self.args = args
+        self.kwargs = kwargs
+
+    def get_credentials(self):
+        if self.supplier:
+            self.username, self.password = self.supplier(*self.args, **self.kwargs)
+        if self.username:
+            auth_args = (self.username, self.password or "")
+        else:
+            auth_args = (self.password,)
+        return auth_args
+
+    def get_password(self, call_supplier=True):
+        if call_supplier and self.supplier:
+            self.username, self.password = self.supplier(*self.args, **self.kwargs)
+        return self.password
+
+    def get_username(self, call_supplier=True):
+        if call_supplier and self.supplier:
+            self.username, self.password = self.supplier(*self.args, **self.kwargs)
+        return self.username
+
+
 class Connection:
     "Manages TCP communication to and from a Redis server"
 
@@ -526,6 +562,7 @@ def __init__(
         username=None,
         retry=None,
         redis_connect_func=None,
+        credentials_provider=None,
     ):
         """
         Initialize a new Connection.
@@ -538,9 +575,10 @@ def __init__(
         self.host = host
         self.port = int(port)
         self.db = db
-        self.username = username
         self.client_name = client_name
-        self.password = password
+        self.credentials_provider = credentials_provider
+        if not self.credentials_provider and (username or password):
+            self.credentials_provider = CredentialsProvider(username, password)
         self.socket_timeout = socket_timeout
         self.socket_connect_timeout = socket_connect_timeout or socket_timeout
         self.socket_keepalive = socket_keepalive
@@ -699,12 +737,9 @@ def on_connect(self):
         "Initialize the connection, authenticate and select a database"
         self._parser.on_connect(self)
 
-        # if username and/or password are set, authenticate
-        if self.username or self.password:
-            if self.username:
-                auth_args = (self.username, self.password or "")
-            else:
-                auth_args = (self.password,)
+        # if credentials provider is set, authenticate
+        if self.credentials_provider:
+            auth_args = self.credentials_provider.get_credentials()
             # avoid checking health here -- PING will fail if we try
             # to check the health prior to the AUTH
             self.send_command("AUTH", *auth_args, check_health=False)
@@ -716,7 +751,11 @@ def on_connect(self):
                 # server seems to be < 6.0.0 which expects a single password
                 # arg. retry auth with just the password.
                 # https://github.com/andymccurdy/redis-py/issues/1274
-                self.send_command("AUTH", self.password, check_health=False)
+                self.send_command(
+                    "AUTH",
+                    self.credentials_provider.get_password(),
+                    check_health=False,
+                )
                 auth_response = self.read_response()
 
             if str_if_bytes(auth_response) != "OK":
@@ -1074,6 +1113,7 @@ def __init__(
         client_name=None,
         retry=None,
         redis_connect_func=None,
+        credentials_provider=None,
     ):
         """
         Initialize a new UnixDomainSocketConnection.
@@ -1085,9 +1125,10 @@ def __init__(
         self.pid = os.getpid()
         self.path = path
         self.db = db
-        self.username = username
         self.client_name = client_name
-        self.password = password
+        self.credentials_provider = credentials_provider
+        if not self.credentials_provider and (username or password):
+            self.credentials_provider = CredentialsProvider(username, password)
         self.socket_timeout = socket_timeout
         self.retry_on_timeout = retry_on_timeout
         if retry_on_error is SENTINEL:

tests/test_connection.py

@@ -1,17 +1,25 @@
+import random
 import socket
+import string
 import types
 from unittest import mock
 from unittest.mock import patch
 
 import pytest
 
+import redis
 from redis.backoff import NoBackoff
-from redis.connection import Connection
-from redis.exceptions import ConnectionError, InvalidResponse, TimeoutError
+from redis.connection import Connection, CredentialsProvider
+from redis.exceptions import (
+    ConnectionError,
+    InvalidResponse,
+    ResponseError,
+    TimeoutError,
+)
 from redis.retry import Retry
 from redis.utils import HIREDIS_AVAILABLE
 
-from .conftest import skip_if_server_version_lt
+from .conftest import _get_client, skip_if_redis_enterprise, skip_if_server_version_lt
 
 
 @pytest.mark.skipif(HIREDIS_AVAILABLE, reason="PythonParser only")
@@ -122,3 +130,86 @@ def test_connect_timeout_error_without_retry(self):
         assert conn._connect.call_count == 1
         assert str(e.value) == "Timeout connecting to server"
         self.clear(conn)
+
+
+class TestCredentialsProvider:
+    @skip_if_redis_enterprise()
+    def test_credentials_provider_without_supplier(self, r, request):
+        # first, test for default user (`username` is supposed to be optional)
+        default_username = "default"
+        temp_pass = "temp_pass"
+        creds_provider = CredentialsProvider(default_username, temp_pass)
+        r.config_set("requirepass", temp_pass)
+        creds = creds_provider.get_credentials()
+        assert r.auth(creds[1], creds[0]) is True
+        assert r.auth(creds_provider.get_password()) is True
+
+        # test for other users
+        username = "redis-py-auth"
+        password = "strong_password"
+
+        def teardown():
+            try:
+                r.auth(temp_pass)
+            except ResponseError:
+                r.auth("default", "")
+            r.config_set("requirepass", "")
+            r.acl_deluser(username)
+
+        request.addfinalizer(teardown)
+
+        assert r.acl_setuser(
+            username,
+            enabled=True,
+            passwords=["+" + password],
+            keys="~*",
+            commands=["+ping", "+command", "+info", "+select", "+flushdb", "+cluster"],
+        )
+
+        creds_provider2 = CredentialsProvider(username, password)
+        r2 = _get_client(
+            redis.Redis, request, flushdb=False, credentials_provider=creds_provider2
+        )
+
+        assert r2.ping() is True
+
+    @skip_if_redis_enterprise()
+    def test_credentials_provider_with_supplier(self, r, request):
+        import functools
+
+        @functools.lru_cache(maxsize=10)
+        def auth_supplier(user, endpoint):
+            def get_random_string(length):
+                letters = string.ascii_lowercase
+                result_str = "".join(random.choice(letters) for i in range(length))
+                return result_str
+
+            auth_token = get_random_string(5) + user + "_" + endpoint
+            return user, auth_token
+
+        username = "redis-py-auth"
+        creds_provider = CredentialsProvider(
+            supplier=auth_supplier,
+            user=username,
+            endpoint="localhost",
+        )
+        password = creds_provider.get_password()
+
+        assert r.acl_setuser(
+            username,
+            enabled=True,
+            passwords=["+" + password],
+            keys="~*",
+            commands=["+ping", "+command", "+info", "+select", "+flushdb", "+cluster"],
+        )
+
+        def teardown():
+            r.acl_deluser(username)
+
+        request.addfinalizer(teardown)
+
+        r2 = _get_client(
+            redis.Redis, request, flushdb=False, credentials_provider=creds_provider
+        )
+
+        assert r2.ping() is True