the new authentication and authorization is tested with hand

src/core/schemas.py

@@ -46,6 +46,9 @@ class UserOutSchema(_UserSchema):
     telegram: str | None = None
     instagram: str | None = None
     twitter: str | None = None
+
+
+class UserRegisterOutSchema(UserOutSchema):
     qr_img: str
 
 

src/core/security.py

@@ -96,8 +96,8 @@ def decode_csrf_token(token) -> CSRFToken:
         return CSRFToken(refresh_token=refresh_token, access_token=access_token)
 
 
-def encode_access_token(username: str, role: UserRolesEnum) -> str:
-    to_encode = {"username": username, "role": role}
+def encode_access_token(username: str, role: UserRolesEnum, refresh_token: str) -> str:
+    to_encode = {"username": username, "role": role, "refresh_token": refresh_token}
     expire = datetime.now(tz=ZoneInfo("UTC")) + timedelta(
         minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES,
     )

src/core/utils.py

@@ -1,5 +1,13 @@
 import asyncio
+import base64
 import functools
+import hashlib
+import io
+
+import qrcode
+from pyotp import totp
+
+from src.core.schemas import UserSchema
 
 
 def asingleton(coro):
@@ -16,3 +24,17 @@ async def wrapper(*args, **kwargs):
         return instance
 
     return wrapper
+
+
+def create_totp_qr_img(user_schema: UserSchema) -> str:
+    provisioning_uri = totp.TOTP(user_schema.totp_hash).provisioning_uri(
+        name=user_schema.username, issuer_name="SimpleRESTBlog"
+    )
+    buffered = io.BytesIO()
+    qrcode.make(provisioning_uri).save(buffered)
+    qr_img = base64.b64encode(buffered.getvalue()).decode()
+    return qr_img
+
+
+def sha256_username(username):
+    return hashlib.sha256(username.encode()).hexdigest()

src/service/user_service.py

@@ -1,8 +1,5 @@
 import asyncio
-import base64
-import io
 
-import qrcode
 from pyotp import random_base32, totp
 
 from src.core.config import settings
@@ -12,10 +9,10 @@
     CredentialsError,
 )
 from src.core.schemas import (
-    UserSchema,
     UserSignupSchema,
     UserOutSchema,
     UserLoginSchema,
+    UserRegisterOutSchema,
 )
 from src.core.security import (
     hash_password,
@@ -25,6 +22,7 @@
     encode_access_token,
     Token,
 )
+from src.core.utils import create_totp_qr_img, sha256_username
 from src.repository.user_repo import UserRepo
 from src.service import Service
 
@@ -48,29 +46,22 @@ async def signup_user(self, user_data: UserSignupSchema) -> UserOutSchema:
         user["totp_hash"] = str(random_base32())
         user_schema = await self.repo.add(user)
 
-        provisioning_uri = totp.TOTP(user_schema.totp_hash).provisioning_uri(
-            name=user_schema.username, issuer_name="SimpleRESTBlog"
-        )
-        buffered = io.BytesIO()
-        qrcode.make(provisioning_uri).save(buffered)
+        qr_img = create_totp_qr_img(user_schema)
 
-        return UserOutSchema(
+        return UserRegisterOutSchema(
             **user_schema.model_dump(),
-            qr_img=base64.b64encode(buffered.getvalue()).decode(),
+            qr_img=qr_img,
         )
 
-    async def authenticate(self, username: str, password: str) -> UserSchema:
-        user = await self.repo.get(username)
-        if verify_password(password, user.password):
-            return user
-        raise CredentialsError()
-
     async def get_user(self, username: str):
         return await self.repo.get(username)
 
     async def login_user(self, user_login: UserLoginSchema) -> Token:
         user = await self.repo.get(user_login.username)
 
+        if not verify_password(user_login.password, user.password):
+            raise CredentialsError()
+
         refresh_token = encode_refresh_token(user.username)
         csrf_token = encode_csrf_token(refresh_token)
 
@@ -85,24 +76,28 @@ async def login_user(self, user_login: UserLoginSchema) -> Token:
             csrf_token=csrf_token,
         )
 
-    async def verify(self, refresh_token: str, username: str, code: str):
-        in_cache_username = await self.redis_client.get(refresh_token, None)
-        if in_cache_username is None or in_cache_username != username:
+    async def _check_refresh_token(self, refresh_token, username):
+        if refresh_token is None:
+            raise CredentialsError("Refresh-Token is not provided")
+        in_cache_username = await self.redis_client.get(refresh_token)
+        if in_cache_username != username:
             raise CredentialsError("Invalid Refresh-Token")
 
+    async def verify(self, refresh_token: str, username: str, code: str):
+        await self._check_refresh_token(refresh_token, username)
         user = await self.repo.get(username)
         if not totp.TOTP(user.totp_hash).verify(code):
             raise CredentialsError("Invalid TOTP code")
 
         await self.redis_client.set(
-            username, True, timeout=settings.TFA_EXPIRE_MINUTES * 60
+            sha256_username(username), True, timeout=settings.TFA_EXPIRE_MINUTES * 60
         )
 
     async def refresh_token(self, old_refresh: str, username: str):
-        in_cache_username, ref_ttl, verified = asyncio.gather(
+        in_cache_username, ref_ttl, verified = await asyncio.gather(
             self.redis_client.get(old_refresh),
             self.redis_client.ttl(old_refresh),
-            self.redis_client.get(username),
+            self.redis_client.get(sha256_username(username)),
         )
         if in_cache_username is None or in_cache_username != username:
             raise CredentialsError("Invalid Refresh-Token")
@@ -112,9 +107,9 @@ async def refresh_token(self, old_refresh: str, username: str):
 
         user = await self.repo.get(username)
 
-        access_token = encode_access_token(username, user.role)
         ref_expire = ref_ttl // 60
         refresh_token = encode_refresh_token(username, ref_expire)
+        access_token = encode_access_token(username, user.role, refresh_token)
         csrf_token = encode_csrf_token(refresh_token, access_token)
 
         await asyncio.gather(
@@ -124,11 +119,13 @@ async def refresh_token(self, old_refresh: str, username: str):
         return Token(access_token, refresh_token, csrf_token)
 
     async def logout(self, refresh_token: str, username: str):
-        if refresh_token is None:
-            raise CredentialsError("Refresh-Token is not provided")
-
-        in_cache_username = await self.redis_client.get(refresh_token)
-        if in_cache_username != username:
-            raise CredentialsError("Invalid Refresh-Token")
+        await self._check_refresh_token(refresh_token, username)
+        await asyncio.gather(
+            self.redis_client.delete(refresh_token),
+            self.redis_client.delete(sha256_username(username)),
+        )
 
-        await self.redis_client.delete(refresh_token)
+    async def get_user_qr_img(self, refresh_token: str, username: str):
+        await self._check_refresh_token(refresh_token, username)
+        user = await self.repo.get(username)
+        return create_totp_qr_img(user)

src/web/auth.py

@@ -42,6 +42,22 @@ async def login(
     response.headers["X-CSRF-TOKEN"] = tokens.csrf_token
 
 
+@router.post("/2fa-img")
+async def get_2fa_image(
+    request: Request,
+    db: Annotated[AsyncSession, Depends(get_db)],
+    redis_client: Annotated[RedisClient, Depends(get_redis_client)],
+    username: Annotated[str, Depends(get_current_username_with_refresh)],
+) -> str:
+    async with UnitOfWork(db):
+        repo = UserRepo(db)
+        service = UserService(repo, redis_client)
+        qr_img = await service.get_user_qr_img(
+            request.cookies.get("Refresh-Token"), username
+        )
+    return qr_img
+
+
 @router.post("/verify")
 async def verify(
     request: Request,