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,