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):
+
+def on_close(exc=None):
     logger.warning(f"RabbitMQ connection closed: {exc}")
 
-async def on_channel_close(ch, exc=None):
+def on_channel_close(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):
+    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(
@@ -159,14 +151,14 @@ async def consume_queue(queue, name: str):
             try:
                 await process_message(message)
             except Exception as e:
-                logger.error(f"Error processing message from {name}: {e}")
+                logging.exception(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))
 
             logger.info(f"Starting message consumption from {name}...")
             await consume_queue(queue, name)
@@ -223,22 +215,21 @@ 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")
-
-    if not gold_url and not golddr_url:
-        raise RuntimeError("No RabbitMQ URLs configured. At least one is required.")
+    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")
 
     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.")
@@ -380,7 +371,7 @@ def watermark(webcam: any, image_data: bytes, tz: str, timestamp: str) -> bytes:
         return buffer.read()
 
     except Exception as e:
-        logger.error(f"Error processing image from camera: {e}")
+        logging.exception(f"Error processing image from camera: {e}")
         return None
 
 def blank_out_image(webcam: any, image_data: bytes, tz: str, timestamp: str) -> bytes:
@@ -426,7 +417,7 @@ def blank_out_image(webcam: any, image_data: bytes, tz: str, timestamp: str) ->
         return buffer.read()
 
     except Exception as e:
-        logger.error(f"Error processing image from camera: {e}")
+        logging.exception(f"Error processing image from camera: {e}")
         return None
 
 
@@ -441,8 +432,8 @@ def save_original_image_to_pvc(camera_id: str, image_bytes: bytes):
         with open(filepath, "wb") as f:
             f.write(image_bytes)
     except Exception as e:
-        logger.error(f"Error saving original image to PVC {filepath}: {e}")
-    logger.info(f"Original image saved to PVC at {filepath}")
+        logging.exception(f"Error saving original image to PVC {filepath}: {e}")
+    logging.info(f"Original image saved to PVC at {filepath}")
 
 def save_watermarked_image_to_pvc(camera_id: str, image_bytes: bytes, timestamp: str, is_on: bool):  
     os.makedirs(os.path.dirname(f'{PVC_WATERMARKED_PATH}'), exist_ok=True)
@@ -456,11 +447,11 @@ def save_watermarked_image_to_pvc(camera_id: str, image_bytes: bytes, timestamp:
         with open(filepath, "wb") as f:
             f.write(image_bytes)
         if is_on:
-            logger.info(f"Watermarked image saved to PVC at {filepath}")
+            logging.info(f"Watermarked image saved to PVC at {filepath}")
         else:
-            logger.info(f"Blank out image saved to PVC at {filepath}")
+            logging.info(f"Blank out image saved to PVC at {filepath}")
     except Exception as e:
-        logger.error(f"Error saving image to PVC {filepath}: {e}")
+        logging.exception(f"Error saving image to PVC {filepath}: {e}")
 
 def save_watermarked_image_to_drivebc_pvc(camera_id: str, image_bytes: bytes, is_on: bool):  
     os.makedirs(os.path.dirname(f'{DRIVEBC_PVC_WATERMARKED_PATH}'), exist_ok=True)
@@ -474,9 +465,9 @@ def save_watermarked_image_to_drivebc_pvc(camera_id: str, image_bytes: bytes, is
         with open(filepath, "wb") as f:
             f.write(image_bytes)
         if is_on:
-            logger.info(f"Watermarked image saved to drivebc PVC at {filepath}")
+            logging.info(f"Watermarked image saved to drivebc PVC at {filepath}")
     except Exception as e:
-        logger.error(f"Error saving image to drivebc PVC {filepath}: {e}")
+        logging.exception(f"Error saving image to drivebc PVC {filepath}: {e}")
 
 def delete_watermarked_image_from_pvc(camera_id: str):
     save_dir = os.path.join(PVC_WATERMARKED_PATH, camera_id)
@@ -491,7 +482,8 @@ def delete_watermarked_image_from_pvc(camera_id: str):
             if os.path.isfile(filepath):
                 os.remove(filepath)
     except Exception as e:
-        logger.error(f"Error deleting watermarked images from PVC {save_dir}: {e}")
+        logging.exception(f"Error deleting watermarked images from PVC {save_dir}: {e}")
+
 
 async def get_images_within(camera_id: str, hours: int = 720) -> list:
     cutoff = datetime.now(timezone.utc) - timedelta(hours=hours)
@@ -754,4 +746,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:
+            logging.exception(f"Failed to refresh RabbitMQ token: {e}")

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

@@ -224,8 +224,8 @@ def setUp(self):
     def tearDown(self):
         super().tearDown()
 
-    # setup_rabbitmq tests
-    @patch("apps.consumer.processor.aio_pika.connect_robust")
+    # # setup_rabbitmq tests
+    @patch("apps.consumer.rabbitmq.RabbitMQTokenConnection.connect", new_callable=AsyncMock)
     def test_setup_rabbitmq_success(self, mock_connect):
         from apps.consumer.processor import setup_rabbitmq
 
@@ -238,14 +238,17 @@ 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")
-            return conn, queue
+            return await setup_rabbitmq(
+                "142.34.229.61",
+                5064
+            )
 
         conn, queue = asyncio.run(run_test())
 
         self.assertEqual(conn, mock_connection)
         self.assertEqual(queue, mock_queue)
-        mock_connect.assert_called_once()
+
+        mock_connect.assert_awaited_once()
 
     # consume_queue tests
     @patch("apps.consumer.processor.process_message", new_callable=AsyncMock)
@@ -323,7 +326,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 +352,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 +381,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 +391,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 +400,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 +418,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 +447,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():