adding oauth configuration

Author: DanielePalaia

poetry.lock

@@ -9,6 +9,8 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = "^3.9"
 python-qpid-proton = "^0.39.0"
+jwt = "^1.3.1"
+pyjwt = "^2.10.1"
 
 [tool.poetry.group.dev.dependencies]
 flake8 = "^7.1.1"

pyproject.toml

@@ -10,6 +10,7 @@
     ExchangeSpecification,
     ExchangeToExchangeBindingSpecification,
     ExchangeToQueueBindingSpecification,
+    OAuth2Options,
     OffsetSpecification,
     RecoveryConfiguration,
     StreamOptions,
@@ -89,4 +90,5 @@
     "Environment",
     "ExchangeCustomSpecification",
     "RecoveryConfiguration",
+    "OAuth2Options",
 ]

rabbitmq_amqp_python_client/__init__.py

@@ -14,7 +14,11 @@
 
 from .address_helper import validate_address
 from .consumer import Consumer
-from .entities import RecoveryConfiguration, StreamOptions
+from .entities import (
+    OAuth2Options,
+    RecoveryConfiguration,
+    StreamOptions,
+)
 from .exceptions import (
     ArgumentOutOfRangeException,
     ValidationCodeException,
@@ -60,6 +64,7 @@ def __init__(
         ssl_context: Union[
             PosixSslConfigurationContext, WinSslConfigurationContext, None
         ] = None,
+        oauth2_options: Optional[OAuth2Options] = None,
         recovery_configuration: RecoveryConfiguration = RecoveryConfiguration(),
     ):
         """
@@ -93,6 +98,7 @@ def __init__(
         self._index: int = -1
         self._publishers: list[Publisher] = []
         self._consumers: list[Consumer] = []
+        self._oauth2_options = oauth2_options
 
         # Some recovery_configuration validation
         if recovery_configuration.back_off_reconnect_interval < timedelta(seconds=1):
@@ -109,19 +115,8 @@ def _set_environment_connection_list(self, connections: []):  # type: ignore
     def _open_connections(self, reconnect_handlers: bool = False) -> None:
 
         logger.debug("inside connection._open_connections creating connection")
-        if self._recovery_configuration.active_recovery is False:
-            self._conn = BlockingConnection(
-                url=self._addr,
-                urls=self._addrs,
-                ssl_domain=self._ssl_domain,
-            )
-        else:
-            self._conn = BlockingConnection(
-                url=self._addr,
-                urls=self._addrs,
-                ssl_domain=self._ssl_domain,
-                on_disconnection_handler=self._on_disconnection,
-            )
+
+        self._create_connection()
 
         if reconnect_handlers is True:
             logger.debug("reconnecting managements, publishers and consumers handlers")
@@ -137,6 +132,35 @@ def _open_connections(self, reconnect_handlers: bool = False) -> None:
                 # Update the broken connection and sender in the consumer
                 self._consumers[i]._update_connection(self._conn)
 
+    def _create_connection(self) -> None:
+
+        user = None
+        password = None
+
+        if self._oauth2_options is not None:
+            user = ""
+            password = self._oauth2_options.token
+
+        if self._recovery_configuration.active_recovery is False:
+            self._conn = BlockingConnection(
+                url=self._addr,
+                urls=self._addrs,
+                oauth2_options=self._oauth2_options,
+                ssl_domain=self._ssl_domain,
+                user=user,
+                password=password,
+            )
+        else:
+            self._conn = BlockingConnection(
+                url=self._addr,
+                urls=self._addrs,
+                oauth2_options=self._oauth2_options,
+                ssl_domain=self._ssl_domain,
+                on_disconnection_handler=self._on_disconnection,
+                user=user,
+                password=password,
+            )
+
     def dial(self) -> None:
         """
         Establish a connection to the AMQP server.

rabbitmq_amqp_python_client/connection.py

@@ -257,3 +257,8 @@ class RecoveryConfiguration:
     active_recovery: bool = True
     back_off_reconnect_interval: timedelta = timedelta(seconds=5)
     MaxReconnectAttempts: int = 5
+
+
+@dataclass
+class OAuth2Options:
+    token: str

rabbitmq_amqp_python_client/entities.py

@@ -9,7 +9,7 @@
 )
 
 from .connection import Connection
-from .entities import RecoveryConfiguration
+from .entities import OAuth2Options, RecoveryConfiguration
 from .ssl_configuration import (
     PosixSslConfigurationContext,
     WinSslConfigurationContext,
@@ -41,6 +41,7 @@ def __init__(
         ssl_context: Union[
             PosixSslConfigurationContext, WinSslConfigurationContext, None
         ] = None,
+        oauth2_options: Optional[OAuth2Options] = None,
         recovery_configuration: RecoveryConfiguration = RecoveryConfiguration(),
     ):
         """
@@ -66,6 +67,7 @@ def __init__(
         self._ssl_context = ssl_context
         self._recovery_configuration = recovery_configuration
         self._connections: list[Connection] = []
+        self._oauth2_options = oauth2_options
 
     def connection(
         self,

rabbitmq_amqp_python_client/environment.py

@@ -1,5 +1,6 @@
 import os
 import sys
+from datetime import datetime, timedelta
 from typing import Optional
 
 import pytest
@@ -9,6 +10,7 @@
     AMQPMessagingHandler,
     Environment,
     Event,
+    OAuth2Options,
     PKCS12Store,
     PosixClientCert,
     PosixSslConfigurationContext,
@@ -22,6 +24,7 @@
 )
 
 from .http_requests import delete_all_connections
+from .utils import token
 
 os.chdir(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
@@ -36,6 +39,20 @@ def environment(pytestconfig):
         environment.close()
 
 
+@pytest.fixture()
+def environment_auth(pytestconfig):
+    token_string = token(datetime.now() + timedelta(milliseconds=2500))
+    environment = Environment(
+        uri="amqp://guest:guest@localhost:5672/",
+        oauth2_options=OAuth2Options(token=token_string),
+    )
+    try:
+        yield environment
+
+    finally:
+        environment.close()
+
+
 @pytest.fixture()
 def connection(pytestconfig):
     environment = Environment(uri="amqp://guest:guest@localhost:5672/")

tests/conftest.py

@@ -43,6 +43,12 @@ def test_connection_ssl(ssl_context) -> None:
     environment.close()
 
 
+def test_connection_auth(environment_auth: Environment) -> None:
+
+    connection = environment_auth.connection()
+    connection.dial()
+
+
 def test_environment_connections_management() -> None:
 
     environment = Environment(uri="amqp://guest:guest@localhost:5672/")

tests/test_connection.py

@@ -1,5 +1,9 @@
+import base64
+from datetime import datetime
 from typing import Optional
 
+import jwt
+
 from rabbitmq_amqp_python_client import (
     AddressHelper,
     Connection,
@@ -92,3 +96,32 @@ def cleanup_dead_lettering(management: Management, bind_path: str) -> None:
     management.unbind(bind_path)
     management.delete_exchange(exchange_dead_lettering)
     management.delete_queue(queue_dead_lettering)
+
+
+def token(duration: datetime) -> str:
+    # Decode the base64 key
+    decoded_key = base64.b64decode("abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGH")
+
+    # Define the claims
+    claims = {
+        "iss": "unit_test",
+        "aud": "rabbitmq",
+        "exp": duration,
+        "scope": ["rabbitmq.configure:*/*", "rabbitmq.write:*/*", "rabbitmq.read:*/*"],
+        "random": random_string(6),
+    }
+
+    # Create the token with the claims and sign it
+    token = jwt.encode(
+        claims, decoded_key, algorithm="HS256", headers={"kid": "token-key"}
+    )
+
+    return token
+
+
+# Helper function to generate a random string (replace with your implementation)
+def random_string(length: int) -> str:
+    import random
+    import string
+
+    return "".join(random.choices(string.ascii_letters + string.digits, k=length))