1 files changed, 56 insertions, 0 deletions

diff --git a/handlers/middleware.py b/handlers/middleware.py
new file mode 100644
index 0000000..87a117a
--- /dev/null
+++ b/handlers/middleware.py
@@ -0,0 +1,56 @@
+from typing import Any, Awaitable, Callable
+
+from aiogram import BaseMiddleware
+from aiogram.types import CallbackQuery, Message, TelegramObject
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from database import sessions
+from models import User
+
+
+class InjectSessionMiddleware(BaseMiddleware):
+    async def __call__[T](
+        self,
+        handler: Callable[[TelegramObject, dict[str, Any]], Awaitable[Any]],
+        event: TelegramObject,
+        data: dict[str, Any],
+    ) -> Any:
+        async with sessions.begin() as session:
+            data["session"] = session
+            handler_result = await handler(event, data)
+        return handler_result
+
+
+class UserAccessMiddleware(BaseMiddleware):
+    async def __call__(
+        self,
+        handler: Callable[[TelegramObject, dict[str, Any]], Awaitable[Any]],
+        event: TelegramObject,
+        data: dict[str, Any],
+    ) -> Any:
+        if not isinstance(event, (Message, CallbackQuery)):
+            raise TypeError(
+                f"UserAccessMiddleware doesn't support event with type: {type(event).__name__}",
+                event,
+            )
+
+        if isinstance(event, Message):
+            user_id = event.chat.id
+        else:
+            user_id = event.from_user.id
+
+        session: AsyncSession = data["session"]
+        user = await session.get(User, user_id)
+        if user is None:
+            error_text = "Вы не добавлены в список пользователей VPN"
+
+            if isinstance(event, Message):
+                await event.answer(error_text)
+            else:
+                await event.answer(error_text)
+
+            return
+
+        data["user"] = user
+
+        return await handler(event, data)