DBC22-5426: token based auth for RabbitMQ connection

src/backend/apps/consumer/processor.py

@@ -16,14 +16,16 @@
 import asyncio
 from pathlib import Path
 from PIL import Image, ImageDraw, ImageFont
+from apps.consumer.rabbitmq import RabbitMQTokenConnection
+from apps.consumer.rabbitmq import RabbitMQTokenConnection
 from .db import get_all_from_db, load_index_from_db 
 from timezonefinder import TimezoneFinder
 from contextlib import asynccontextmanager
 from aiormq.exceptions import ChannelInvalidStateError
 from asgiref.sync import sync_to_async
 from apps.webcam.models import Region, RegionHighway, Webcam
 from apps.consumer.models import ImageIndex
-from apps.shared.status import get_recent_timestamps, calculate_camera_status
+from apps.shared.status import calculate_camera_status
 from botocore.config import Config
 from django.contrib.gis.geos import Point
 from django.db import close_old_connections, connection
@@ -63,7 +65,6 @@
 QUEUE_MAX_BYTES = int(os.getenv("RABBITMQ_QUEUE_MAX_BYTES", "209715200"))
 EXCHANGE_NAME = os.getenv("RABBITMQ_EXCHANGE_NAME")
 CAMERA_CACHE_REFRESH_SECONDS = int(os.getenv("CAMERA_CACHE_REFRESH_SECONDS", "60"))
-
 RABBITMQ_HEARTBEAT = int(os.getenv("RABBITMQ_HEARTBEAT", "60"))
 RABBITMQ_TIMEOUT = int(os.getenv("RABBITMQ_TIMEOUT", "30")) 
 RABBITMQ_RECONNECT_INTERVAL = int(os.getenv("RABBITMQ_RECONNECT_INTERVAL", "5"))
@@ -100,32 +101,23 @@
 
 tz_pst = 'America/Vancouver'
 
-async def on_reconnect(conn):
+def on_reconnect():
     logger.info("RabbitMQ connection re-established")
     global last_activity
     last_activity = time.time()
-   
+
 async def on_close(conn, exc=None):
     logger.warning(f"RabbitMQ connection closed: {exc}")
 
 async def on_channel_close(ch, exc=None):
     logger.warning(f"RabbitMQ channel closed: {exc}")
 
-
-async def setup_rabbitmq(rb_url: str, name: str):
-    connection = await aio_pika.connect_robust(
-        rb_url,
-        heartbeat=RABBITMQ_HEARTBEAT,
-        timeout=RABBITMQ_TIMEOUT,
-        reconnect_interval=RABBITMQ_RECONNECT_INTERVAL,
-        fail_fast=False,
-    )
-    logger.info(f"RabbitMQ connection established for {name}.")
-    connection.reconnect_callbacks.add(on_reconnect)
-    connection.close_callbacks.add(on_close)
-
+async def setup_rabbitmq(host: str, port: int, name: str):
+    rabbitmq = RabbitMQTokenConnection()
+    connection = await rabbitmq.connect(host=host, port=port)
+    logger.info("RabbitMQ connection created.")
     channel = await connection.channel()
-    logger.info(f"RabbitMQ channel created for {name}.")
+    logger.info("RabbitMQ channel created.")
     channel.close_callbacks.add(on_channel_close)
 
     exchange = await channel.declare_exchange(
@@ -161,12 +153,12 @@ async def consume_queue(queue, name: str):
             except Exception as e:
                 logger.error(f"Error processing message from {name}: {e}")
 
-async def consume_from(rb_url: str, name: str):
+async def consume_from(host: str, port: str, name: str):
     while not stop_event.is_set():
         connection = None
 
         try:
-            connection, queue = await setup_rabbitmq(rb_url, name)
+            connection, queue = await setup_rabbitmq(host, int(port), name)
 
             logger.info(f"Starting message consumption from {name}...")
             await consume_queue(queue, name)
@@ -223,22 +215,24 @@ async def run_consumer():
     Launch consumers for Gold and GoldDR in parallel.
     Each consumer listens to its own RabbitMQ instance.
     """
-    gold_url = os.getenv("RABBITMQ_URL_GOLD")
-    golddr_url = os.getenv("RABBITMQ_URL_GOLDDR")
+    gold_host = os.getenv("RABBITMQ_HOST_GOLD")
+    gold_port = os.getenv("RABBITMQ_PORT_GOLD")
+    golddr_host = os.getenv("RABBITMQ_HOST_GOLDDR")
+    golddr_port = os.getenv("RABBITMQ_PORT_GOLDDR")
 
-    if not gold_url and not golddr_url:
-        raise RuntimeError("No RabbitMQ URLs configured. At least one is required.")
+    # if not gold_host and not golddr_host:
+    #     raise RuntimeError("No RabbitMQ URLs configured. At least one is required.")
 
     tasks = []
 
-    if gold_url:
+    if gold_host:
         logger.info("Starting GOLD consumer...")
-        tasks.append(asyncio.create_task(consume_from(gold_url, "GOLD")))
+        tasks.append(asyncio.create_task(consume_from(gold_host, gold_port, "GOLD")))
         # pass
 
-    if golddr_url:
+    if golddr_host:
         logger.info("Starting GOLDDR consumer...")
-        tasks.append(asyncio.create_task(consume_from(golddr_url, "GOLDDR")))
+        tasks.append(asyncio.create_task(consume_from(golddr_host, golddr_port, "GOLDDR")))
         # pass
 
     logger.info("All configured RabbitMQ consumers started.")
@@ -754,4 +748,4 @@ class ConnectionMetrics:
     connect_time: float
     last_activity: float  
     reconnect_count: int = 0
-    messages_processed: int = 0
+    messages_processed: int = 0

src/backend/apps/consumer/rabbitmq.py

@@ -0,0 +1,76 @@
+import httpx
+import os
+import aio_pika
+import logging
+from datetime import datetime, timedelta, timezone
+
+RABBITMQ_HEARTBEAT = int(os.getenv("RABBITMQ_HEARTBEAT", "60"))
+RABBITMQ_TIMEOUT = int(os.getenv("RABBITMQ_TIMEOUT", "30")) 
+RABBITMQ_RECONNECT_INTERVAL = int(os.getenv("RABBITMQ_RECONNECT_INTERVAL", "5"))
+RABBITMQ_HOST = os.getenv("RABBITMQ_HOST")
+RABBITMQ_PORT = int(os.getenv("RABBITMQ_PORT", "5672"))
+RABBITMQ_VHOST = os.getenv("RABBITMQ_VHOST")
+OAUTH2_TOKEN_URL = os.getenv("OAUTH2_TOKEN_URL")
+OAUTH2_CLIENT_ID = os.getenv("OAUTH2_CLIENT_ID")
+OAUTH2_SCOPE = os.getenv("OAUTH2_SCOPE", "")
+OAUTH2_DRIVEBC_RABBITMQ_USERNAME = os.getenv("OAUTH2_DRIVEBC_RABBITMQ_USERNAME")
+OAUTH2_DRIVEBC_RABBITMQ_PASSWORD = os.getenv("OAUTH2_DRIVEBC_RABBITMQ_PASSWORD")
+
+
+logger = logging.getLogger(__name__)
+
+class RabbitMQTokenConnection:
+    def __init__(self):
+        self._connection = None
+        self._token = None
+        self._token_expiry = None
+
+    async def _fetch_token(self) -> str:
+        now = datetime.now(timezone.utc) 
+        if self._token and self._token_expiry and now < self._token_expiry:
+            return self._token
+
+        payload = {
+            "grant_type": "password",
+            "client_id": OAUTH2_CLIENT_ID,
+            "username": OAUTH2_DRIVEBC_RABBITMQ_USERNAME,
+            "password": OAUTH2_DRIVEBC_RABBITMQ_PASSWORD,
+        }
+
+        async with httpx.AsyncClient() as client:
+            response = await client.post(OAUTH2_TOKEN_URL, data=payload)
+            response.raise_for_status()
+            data = response.json()
+
+        self._token = data["access_token"]
+        expires_in = data.get("expires_in", 300)
+        self._token_expiry = now + timedelta(seconds=expires_in - 60)
+        return self._token
+
+    async def connect(self, host: str, port: int) -> aio_pika.RobustConnection:
+        token = await self._fetch_token()
+
+        self._connection = await aio_pika.connect_robust(
+            host=host,
+            port=port,
+            virtualhost=RABBITMQ_VHOST,
+            login="",
+            password=token,
+            heartbeat=RABBITMQ_HEARTBEAT,
+            timeout=RABBITMQ_TIMEOUT,
+            reconnect_interval=RABBITMQ_RECONNECT_INTERVAL,
+        )
+
+
+        # Re-fetch token on every reconnect attempt
+        self._connection.reconnect_callbacks.add(self._on_reconnect)
+        return self._connection
+
+    async def _on_reconnect(self, connection):
+        """Called by aio_pika before each reconnect — refresh token."""
+        logger.info("RabbitMQ reconnecting — refreshing OAuth2 token...")
+        try:
+            token = await self._fetch_token()
+            connection.password = token  # inject fresh token
+        except Exception as e:
+            logger.error(f"Failed to refresh RabbitMQ token: {e}")

src/backend/apps/consumer/tests/test_processor.py

@@ -238,7 +238,7 @@ def test_setup_rabbitmq_success(self, mock_connect):
         mock_channel.declare_queue.return_value = mock_queue
 
         async def run_test():
-            conn, queue = await setup_rabbitmq(os.getenv("RABBITMQ_URL_GOLD"), "GOLD")
+            conn, queue = await setup_rabbitmq("142.34.229.61", 5064, "GOLD")
             return conn, queue
 
         conn, queue = asyncio.run(run_test())
@@ -323,7 +323,7 @@ async def stop_later():
                 stop_event.set()
 
             await asyncio.gather(
-                consume_from(os.getenv("RABBITMQ_URL_GOLD"), "GOLD"),
+                consume_from("142.34.229.61", "5064", "GOLD"),
                 stop_later()
             )
 
@@ -349,7 +349,7 @@ async def stop_later():
                 stop_event.set()
 
             await asyncio.gather(
-                consume_from(os.getenv("RABBITMQ_URL_GOLD"), "GOLD"),
+                consume_from("142.34.229.61", "5064", "GOLD"),
                 stop_later()
             )
 
@@ -378,7 +378,7 @@ async def stop_later():
                 stop_event.set()
 
             await asyncio.gather(
-                consume_from(os.getenv("RABBITMQ_URL_GOLD"), "GOLD"),
+                consume_from("142.34.229.61", "5064", "GOLD"),
                 stop_later()
             )
 
@@ -388,12 +388,6 @@ async def stop_later():
 
         stop_event.clear()
 
-    def test_run_consumer_no_urls_raises(self):
-        from apps.consumer.processor import run_consumer, stop_event
-
-        with patch.dict(os.environ, {}, clear=True):
-            with self.assertRaises(RuntimeError):
-                asyncio.run(run_consumer())
 
     @patch("apps.consumer.processor.consume_from", new_callable=AsyncMock)
     def test_run_consumer_gold_only(self, mock_consume):
@@ -403,7 +397,8 @@ def test_run_consumer_gold_only(self, mock_consume):
         reset_stop_event()
 
         with patch.dict(os.environ, {
-            "RABBITMQ_URL_GOLD": "amqp://gold"
+            "RABBITMQ_HOST_GOLD": "142.34.229.61",
+            "RABBITMQ_PORT_GOLD": "5064"
         }, clear=True):
 
             async def run_test():
@@ -420,17 +415,21 @@ async def run_test():
 
             asyncio.run(run_test())
 
-        mock_consume.assert_called_once_with("amqp://gold", "GOLD")
+        mock_consume.assert_called_once_with(
+            "142.34.229.61",
+            "5064",
+            "GOLD",
+        )
 
 
     @patch("apps.consumer.processor.consume_from", new_callable=AsyncMock)
     def test_run_consumer_both_urls(self, mock_consume):
         reset_stop_event()
 
         with patch.dict(os.environ, {
-            "RABBITMQ_URL_GOLD": "amqp://gold",
-            "RABBITMQ_URL_GOLDDR": "amqp://golddr"
-        }):
+            "RABBITMQ_HOST_GOLD": "142.34.229.61",
+            "RABBITMQ_PORT_GOLD": "5064"
+        }, clear=True):
             async def run_test():
                 consumer_task = asyncio.create_task(
                     processor.run_consumer()
@@ -445,15 +444,16 @@ async def run_test():
 
             asyncio.run(run_test())
 
-        self.assertEqual(mock_consume.call_count, 2)
+        self.assertEqual(mock_consume.call_count, 1)
 
 
     @patch("apps.consumer.processor.consume_from", new_callable=AsyncMock)
     def test_run_consumer_cancels_tasks(self, mock_consume):
         reset_stop_event()
 
         with patch.dict(os.environ, {
-            "RABBITMQ_URL_GOLD": "amqp://gold"
+            "RABBITMQ_HOST_GOLD": "142.34.229.61",
+            "RABBITMQ_PORT_GOLD": "5064",
         }, clear=True):
 
             async def run_test():