Skip to content

Commit efae570

Browse files
CaptainArilCaptainAril
committed
inter-service broker security
1 parent 341d3ee commit efae570

File tree

11 files changed

+153
-48
lines changed

11 files changed

+153
-48
lines changed

‎src/api/constants/queues.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
from typing import TypedDict
2+
3+
4+
class QueueNames(TypedDict):
5+
USER_REGISTRATION: str
6+
EMAIL_VALIDATION: str
7+
8+
9+
QUEUE_NAMES: QueueNames = {
10+
"USER_REGISTRATION": "create_user",
11+
"EMAIL_VALIDATION": "validate_email",
12+
}

‎src/api/constants/signature_sources.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
from typing import TypedDict
2+
3+
4+
class SignatureSources(TypedDict):
5+
gateway: str
6+
queue: str
7+
8+
9+
SIGNATURE_SOURCES: SignatureSources = {
10+
"gateway": "Gateway",
11+
"queue": "Broker Queue",
12+
}

‎src/api/middlewares/BrokerMiddleware.py‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
from types import CoroutineType
2+
from typing import Any
3+
from collections.abc import Callable, Awaitable
4+
5+
from faststream import BaseMiddleware
6+
from faststream.broker.message import StreamMessage
7+
from faststream.rabbit.message import RabbitMessage
8+
9+
from src.env import queue
10+
from src.utils.logger import Logger
11+
from src.api.services.UtilityService import UtilityService
12+
from src.api.constants.signature_sources import SIGNATURE_SOURCES
13+
14+
15+
class PublishMiddleware(BaseMiddleware):
16+
"""
17+
Middleware to handle subscription messages.
18+
"""
19+
20+
async def publish_scope(
21+
self,
22+
call_next: Callable[..., Awaitable[Any]],
23+
msg: RabbitMessage,
24+
*args: tuple[Any, ...],
25+
**kwargs: dict[str, Any],
26+
) -> CoroutineType:
27+
timestamp = UtilityService.get_timestamp()
28+
signature = UtilityService.generate_signature(queue["key"], timestamp)
29+
30+
headers = {}
31+
32+
headers["X-BROKER-SIGNATURE"] = signature
33+
headers["X-BROKER-TIMESTAMP"] = timestamp
34+
headers["X-BROKER-KEY"] = queue["key"]
35+
36+
kwargs["headers"] = headers
37+
return await super().publish_scope(call_next, msg, *args, **kwargs)
38+
39+
40+
class SubscribeMiddleware(BaseMiddleware):
41+
async def consume_scope(
42+
self, call_next: Callable[[Any], Awaitable[Any]], msg: StreamMessage
43+
) -> CoroutineType | None:
44+
logger = Logger(__name__)
45+
try:
46+
UtilityService.verify_signature(
47+
logger=logger,
48+
signature_data={
49+
"signature": msg.headers["X-BROKER-SIGNATURE"],
50+
"timestamp": msg.headers["X-BROKER-TIMESTAMP"],
51+
"key": queue["key"],
52+
"ttl": queue["ttl"],
53+
"title": SIGNATURE_SOURCES["gateway"],
54+
},
55+
)
56+
57+
return await super().consume_scope(call_next, msg)
58+
except KeyError as e:
59+
message = f"Missing required header: {e}"
60+
logger.error(
61+
{
62+
"activity_type": "Authenticate GatewaBroker Queue Request",
63+
"message": message,
64+
"metadata": {"headers": msg.headers},
65+
}
66+
)
67+
except Exception as e:
68+
queue_operation = msg.raw_message.routing_key
69+
message = f"`{queue_operation}` operation failed: {e}"
70+
logger.error(
71+
{
72+
"activity_type": "Authenticate GatewaBroker Queue Request",
73+
"message": message,
74+
"metadata": {"headers": msg.headers, "message": msg._decoded_body},
75+
}
76+
)

‎src/api/middlewares/GateWayMiddleware.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
from django.http import HttpRequest
22
from ninja.errors import AuthenticationError
3-
from ninja.openapi.schema import OpenAPISchema
43
from ninja.security import APIKeyHeader
4+
from ninja.openapi.schema import OpenAPISchema
55

6-
from src.api.constants.signature_sources import SIGNATURE_SOURCES
7-
from src.api.services.UtilityService import SignatureData, UtilityService
86
from src.env import api_gateway
97
from src.utils.logger import Logger
8+
from src.api.services.UtilityService import SignatureData, UtilityService
9+
from src.api.constants.signature_sources import SIGNATURE_SOURCES
1010

1111

1212
class GateWayAuth(APIKeyHeader):
@@ -43,7 +43,7 @@ def authenticate(self, request: HttpRequest, key: str | None) -> str | None:
4343
}
4444
)
4545
raise AuthenticationError(message=message)
46-
46+
4747
signature_data: SignatureData = {
4848
"signature": api_signature,
4949
"timestamp": api_timestamp,
@@ -52,7 +52,9 @@ def authenticate(self, request: HttpRequest, key: str | None) -> str | None:
5252
"title": SIGNATURE_SOURCES["gateway"],
5353
}
5454

55-
UtilityService.verify_signature(logger=self.logger, signature_data=signature_data)
55+
UtilityService.verify_signature(
56+
logger=self.logger, signature_data=signature_data
57+
)
5658

5759
self.logger.debug(
5860
{

‎src/api/services/UserService.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,16 @@
22

33
from faststream.rabbit import RabbitRouter
44

5-
from src.api.constants.activity_types import ACTIVITY_TYPES
6-
from src.api.constants.messages import DYNAMIC_MESSAGES, MESSAGES
5+
from src.utils.svcs import Service
6+
from src.utils.logger import Logger
7+
from src.api.models.postgres import User
78
from src.api.constants.queues import QUEUE_NAMES
9+
from src.api.constants.messages import MESSAGES, DYNAMIC_MESSAGES
10+
from src.api.typing.UserSuccess import UserSuccess
11+
from src.api.constants.activity_types import ACTIVITY_TYPES
12+
from src.api.repositories.UserRepository import UserRepository
813
from src.api.models.payload.requests.CreateUserRequest import CreateUserRequest
914
from src.api.models.payload.requests.UpdateUserRequest import UpdateUserRequest
10-
from src.api.models.postgres import User
11-
from src.api.repositories.UserRepository import UserRepository
12-
from src.api.typing.UserSuccess import UserSuccess
13-
from src.utils.logger import Logger
14-
from src.utils.svcs import Service
1515

1616
from .UtilityService import UtilityService
1717

‎src/api/services/UtilityService.py‎

Lines changed: 15 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,34 @@
1-
import hashlib
21
import hmac
3-
from datetime import datetime, timedelta
4-
from typing import TypedDict
2+
import hashlib
53
from uuid import uuid4
4+
from typing import TypedDict
5+
from datetime import datetime, timedelta
66

7-
import bcrypt
87
import jwt
9-
from django.utils import timezone
8+
import bcrypt
109
from faker import Faker
10+
from django.utils import timezone
1111
from ninja.errors import AuthenticationError
1212

13-
from src.api.enums.CharacterCasing import CharacterCasing
14-
from src.api.models.postgres import User
15-
from src.api.typing.ExpireUUID import ExpireUUID
1613
from src.env import jwt_config
17-
from src.utils.logger import Logger
1814
from src.utils.svcs import Service
15+
from src.utils.logger import Logger
16+
from src.api.models.postgres import User
17+
from src.api.typing.ExpireUUID import ExpireUUID
18+
from src.api.enums.CharacterCasing import CharacterCasing
1919

2020
DEFAULT_CHARACTER_LENGTH = 12
2121
fake = Faker()
2222

23+
2324
class SignatureData(TypedDict):
2425
title: str
2526
signature: str
2627
timestamp: str
2728
key: str
2829
ttl: int | float
2930

31+
3032
@Service()
3133
class UtilityService:
3234
@staticmethod
@@ -91,25 +93,22 @@ def generate_uuid() -> ExpireUUID:
9193
lifespan = timedelta(hours=24)
9294
expires_at = current_time + lifespan
9395
return {"uuid": uuid4(), "expires_at": expires_at}
94-
96+
9597
@staticmethod
9698
def generate_signature(key: str, timestamp: str) -> str:
9799
signature = hmac.new(
98100
key=key.encode(), msg=timestamp.encode(), digestmod=hashlib.sha256
99101
).hexdigest()
100102
return signature
101103

102-
103104
@staticmethod
104105
def verify_signature(signature_data: SignatureData, logger: Logger) -> bool:
105-
106106
signature = signature_data["signature"]
107107
timestamp = signature_data["timestamp"]
108108
key = signature_data["key"]
109109
ttl = signature_data["ttl"]
110110
title = signature_data["title"]
111111

112-
113112
valid_signature = UtilityService.generate_signature(key, timestamp)
114113
is_valid = hmac.compare_digest(valid_signature, signature)
115114

@@ -124,9 +123,9 @@ def verify_signature(signature_data: SignatureData, logger: Logger) -> bool:
124123
)
125124
raise AuthenticationError(message=message)
126125

127-
initial_time = datetime.fromtimestamp(float(timestamp)/ 1000)
126+
initial_time = datetime.fromtimestamp(float(timestamp) / 1000)
128127
valid_window = initial_time + timedelta(minutes=ttl)
129-
128+
130129
if valid_window < datetime.now():
131130
message = "Signature expired!"
132131
logger.error(
@@ -140,8 +139,7 @@ def verify_signature(signature_data: SignatureData, logger: Logger) -> bool:
140139

141140
return True
142141

143-
144142
@staticmethod
145143
def get_timestamp() -> str:
146144
current_time = datetime.now().timestamp() * 1000
147-
return str(current_time)
145+
return str(current_time)

‎src/config/asgi.py‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@
1111

1212
from django.core.asgi import get_asgi_application
1313
from faststream.rabbit import RabbitBroker
14-
from starlette.applications import Starlette
1514
from starlette.routing import Mount
15+
from starlette.applications import Starlette
1616

1717
from src.env import rabbitmq_config
1818

@@ -21,9 +21,12 @@
2121

2222
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "src.config.settings")
2323

24-
def setup_broker_middlewares():
25-
from src.api.middlewares.BrokerMiddleware import (PublishMiddleware,
26-
SubscribeMiddleware)
24+
25+
def setup_broker_middlewares() -> None:
26+
from src.api.middlewares.BrokerMiddleware import (
27+
PublishMiddleware,
28+
SubscribeMiddleware,
29+
)
2730

2831
broker.add_middleware(SubscribeMiddleware)
2932
broker.add_middleware(PublishMiddleware)
@@ -36,5 +39,4 @@ def setup_broker_middlewares():
3639
)
3740

3841

39-
from src.api.services.external import \
40-
RabbitMQRoutes as RabbitMQRoutes # noqa: E402
42+
from src.api.services.external import RabbitMQRoutes as RabbitMQRoutes # noqa: E402

‎src/config/settings.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,11 @@
22

33
from .apps import INSTALLED_APPS as INSTALLED_APPS
44
from .caches import CACHES as CACHES
5-
from .databases import DATABASE_ROUTERS as DATABASE_ROUTERS
6-
from .databases import DATABASES as DATABASES
75
from .logger import LOGGING as LOGGING
8-
from .middleware import MIDDLEWARE as MIDDLEWARE
6+
from .databases import DATABASES as DATABASES
7+
from .databases import DATABASE_ROUTERS as DATABASE_ROUTERS
98
from .templates import TEMPLATES as TEMPLATES
9+
from .middleware import MIDDLEWARE as MIDDLEWARE
1010

1111
SECRET_KEY = app["secret_key"]
1212
DEBUG = app["debug"]
@@ -20,10 +20,10 @@
2020
MEDIA_URL = "media/"
2121

2222

23-
LANGUAGE_CODE = 'en-us'
23+
LANGUAGE_CODE = "en-us"
2424

25-
TIME_ZONE = 'UTC'
25+
TIME_ZONE = "UTC"
2626

2727
USE_I18N = True
2828

29-
USE_TZ = True
29+
USE_TZ = True

‎src/env.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from typing import TypedDict
22

3-
from src import __description__, __display_name__, __name__, __version__
4-
from src.utils.env import get_env_float, get_env_int, get_env_list, get_env_str
3+
from src import __name__, __version__, __description__, __display_name__
4+
from src.utils.env import get_env_int, get_env_str, get_env_list, get_env_float
55

66

77
class Env:
@@ -58,10 +58,12 @@ class Gateway(TypedDict):
5858
key: str
5959
ttl: int
6060

61+
6162
class Queue(TypedDict):
6263
key: str
6364
ttl: float
6465

66+
6567
env = Env()
6668

6769
app: App = {
@@ -134,6 +136,6 @@ class Queue(TypedDict):
134136
"jwt_config",
135137
"log",
136138
"otp",
137-
"rabbitmq_config",
138139
"queue",
140+
"rabbitmq_config",
139141
]

‎src/utils/env/__init__.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1-
from .env import get_env_float, get_env_int, get_env_list, get_env_str
1+
from .env import get_env_int, get_env_str, get_env_list, get_env_float
22

33
__all__ = [
4+
"get_env_float",
45
"get_env_int",
56
"get_env_list",
67
"get_env_str",
7-
"get_env_float",
88
]

0 commit comments

Comments
 (0)