GeezerVison

The World Wide Web, restored to 800 pixels of glory.

GeezerVison is showing github.com — feat: 2fa implementation · IDTS-LAB/python-codebase@98646e6

File tree

Expand file treeCollapse file tree

Open diff view settings

Filter options

src
- core/security
  - two_factor_auth.py
- modules/user
  - presentation
    - dependency.py
    - routers
      - two_factor_router.py
  - providers.py

Expand file treeCollapse file tree

Open diff view settings

Collapse file

`‎src/core/security/two_factor_auth.py‎`

Copy file name to clipboardExpand all lines: src/core/security/two_factor_auth.py

+41-39Lines changed: 41 additions & 39 deletions

Original file line number	Diff line number	Diff line change
`@@ -2,7 +2,6 @@`
`2`	`2`
`3`	`3`	`from __future__ import annotations`
`4`	`4`
`5`		`-import base64`
`6`	`5`	`import secrets`
`7`	`6`	`from typing import Literal`
`8`	`7`	`from uuid import UUID`
`@@ -11,14 +10,14 @@`
`11`	`10`
`12`	`11`	`from src.core.config.setting import get_settings`
`13`	`12`	`from src.core.email.service import EmailService`
`14`		`-from src.modules.user.domain.entities.user import User, UserSecurity`
	`13`	`+from src.modules.user.domain.entities.user import User`
`15`	`14`	`from src.modules.user.domain.repositories.user_repository import UserRepository`
`16`	`15`	`from src.shared.unit_of_work import UnitOfWork`
`17`	`16`
`18`	`17`
`19`	`18`	`class TwoFactorAuthService:`
`20`	`19`	`"""Service for managing two-factor authentication.`
`21`		`-`
	`20`	`+`
`22`	`21`	`Supports:`
`23`	`22`	`- TOTP (Time-based One-Time Password) for authenticator apps like Google Authenticator, Authy, etc.`
`24`	`23`	`- Email-based 2FA codes`
`@@ -37,7 +36,7 @@ def __init__(`
`37`	`36`
`38`	`37`	`async def setup_totp(self, user_id: UUID) -> dict[str, str]:`
`39`	`38`	`"""Set up TOTP for a user.`
`40`		`-`
	`39`	`+`
`41`	`40`	`Returns:`
`42`	`41`	`dict with 'secret', 'uri', and 'qr_code_data' keys`
`43`	`42`	`"""`
`@@ -48,14 +47,13 @@ async def setup_totp(self, user_id: UUID) -> dict[str, str]:`
`48`	`47`
`49`	`48`	`# Generate a new secret`
`50`	`49`	`secret = pyotp.random_base32()`
`51`		`-`
	`50`	`+`
`52`	`51`	`# Create TOTP URI for QR code generation`
`53`	`52`	`issuer = self._settings.JWT_ISSUER or "TodoApp"`
`54`	`53`	`uri = pyotp.totp.TOTP(secret).provisioning_uri(`
`55`		`- name=user.email,`
`56`		`- issuer_name=issuer`
	`54`	`+ name=user.email, issuer_name=issuer`
`57`	`55`	`)`
`58`		`-`
	`56`	`+`
`59`	`57`	`# Store the secret temporarily (not enabled yet)`
`60`	`58`	`user.security.two_factor_secret = secret`
`61`	`59`	`user.security.two_factor_enabled = False`
`@@ -65,16 +63,16 @@ async def setup_totp(self, user_id: UUID) -> dict[str, str]:`
`65`	`63`	`return {`
`66`	`64`	`"secret": secret,`
`67`	`65`	`"uri": uri,`
`68`		`- "qr_code_data": f"otpauth://totp/{issuer}:{user.email}?secret={secret}&issuer={issuer}"`
	`66`	`+ "qr_code_data": f"otpauth://totp/{issuer}:{user.email}?secret={secret}&issuer={issuer}",`
`69`	`67`	`}`
`70`	`68`
`71`	`69`	`async def verify_totp_setup(self, user_id: UUID, code: str) -> dict[str, list[str]]:`
`72`	`70`	`"""Verify TOTP setup and enable 2FA.`
`73`		`-`
	`71`	`+`
`74`	`72`	`Args:`
`75`	`73`	`user_id: The user's ID`
`76`	`74`	`code: The TOTP code from the authenticator app`
`77`		`-`
	`75`	`+`
`78`	`76`	`Returns:`
`79`	`77`	`dict with 'backup_codes' key containing recovery codes`
`80`	`78`	`"""`
`@@ -93,7 +91,7 @@ async def verify_totp_setup(self, user_id: UUID, code: str) -> dict[str, list[st`
`93`	`91`
`94`	`92`	`# Generate backup codes`
`95`	`93`	`backup_codes = [secrets.token_hex(4) for _ in range(10)]`
`96`		`-`
	`94`	`+`
`97`	`95`	`# Enable 2FA and store backup codes`
`98`	`96`	`user.security.two_factor_enabled = True`
`99`	`97`	`user.security.two_factor_backup_codes = ",".join(backup_codes)`
`@@ -104,11 +102,11 @@ async def verify_totp_setup(self, user_id: UUID, code: str) -> dict[str, list[st`
`104`	`102`
`105`	`103`	`async def disable_totp(self, user_id: UUID, code: str) -> bool:`
`106`	`104`	`"""Disable TOTP 2FA for a user.`
`107`		`-`
	`105`	`+`
`108`	`106`	`Args:`
`109`	`107`	`user_id: The user's ID`
`110`	`108`	`code: Current TOTP code or backup code for verification`
`111`		`-`
	`109`	`+`
`112`	`110`	`Returns:`
`113`	`111`	`True if successfully disabled`
`114`	`112`	`"""`
`@@ -122,7 +120,7 @@ async def disable_totp(self, user_id: UUID, code: str) -> bool:`
`122`	`120`
`123`	`121`	`# Verify code`
`124`	`122`	`verified = False`
`125`		`-`
	`123`	`+`
`126`	`124`	`# Check if it's a backup code`
`127`	`125`	`if user.security.two_factor_backup_codes:`
`128`	`126`	`backup_codes = user.security.two_factor_backup_codes.split(",")`
`@@ -151,10 +149,10 @@ async def disable_totp(self, user_id: UUID, code: str) -> bool:`
`151`	`149`
`152`	`150`	`async def send_email_2fa_code(self, user_id: UUID) -> bool:`
`153`	`151`	`"""Send a 2FA code via email.`
`154`		`-`
	`152`	`+`
`155`	`153`	`Args:`
`156`	`154`	`user_id: The user's ID`
`157`		`-`
	`155`	`+`
`158`	`156`	`Returns:`
`159`	`157`	`True if email was sent successfully`
`160`	`158`	`"""`
`@@ -168,14 +166,15 @@ async def send_email_2fa_code(self, user_id: UUID) -> bool:`
`168`	`166`
`169`	`167`	`# Generate a 6-digit code`
`170`	`168`	`code = secrets.token_hex(3)[:6]`
`171`		`-`
	`169`	`+`
`172`	`170`	`# Store the code temporarily in security settings (with expiry info)`
`173`	`171`	`# In production, you'd want to store this in Redis with TTL`
`174`	`172`	`if not user.security:`
`175`	`173`	`raise ValueError("User security not found")`
`176`		`-`
	`174`	`+`
`177`	`175`	`# Store code with timestamp (format: "code:timestamp")`
`178`	`176`	`import time`
	`177`	`+`
`179`	`178`	`user.security.two_factor_secret = f"email_code:{code}:{int(time.time())}"`
`180`	`179`	`await self._user_repository.save_security(user.security)`
`181`	`180`	`await self._unit_of_work.commit()`
`@@ -191,22 +190,22 @@ async def send_email_2fa_code(self, user_id: UUID) -> bool:`
`191`	`190`	`</body>`
`192`	`191`	`</html>`
`193`	`192`	`"""`
`194`		`-`
	`193`	`+`
`195`	`194`	`await self._email_service.send_email(`
`196`	`195`	`to=user.email,`
`197`	`196`	`subject="Your Verification Code",`
`198`	`197`	`html_body=html_body,`
`199`	`198`	`)`
`200`		`-`
	`199`	`+`
`201`	`200`	`return True`
`202`	`201`
`203`	`202`	`async def verify_email_2fa_code(self, user_id: UUID, code: str) -> bool:`
`204`	`203`	`"""Verify an email-based 2FA code.`
`205`		`-`
	`204`	`+`
`206`	`205`	`Args:`
`207`	`206`	`user_id: The user's ID`
`208`	`207`	`code: The code received via email`
`209`		`-`
	`208`	`+`
`210`	`209`	`Returns:`
`211`	`210`	`True if code is valid`
`212`	`211`	`"""`
`@@ -222,17 +221,18 @@ async def verify_email_2fa_code(self, user_id: UUID, code: str) -> bool:`
`222`	`221`	`parts = stored.split(":")`
`223`	`222`	`if len(parts) != 3:`
`224`	`223`	`raise ValueError("Invalid code format")`
`225`		`-`
	`224`	`+`
`226`	`225`	`stored_code = parts[1]`
`227`	`226`	`timestamp = int(parts[2])`
`228`		`-`
	`227`	`+`
`229`	`228`	`import time`
	`229`	`+`
`230`	`230`	`current_time = int(time.time())`
`231`		`-`
	`231`	`+`
`232`	`232`	`# Code expires after 10 minutes`
`233`	`233`	`if current_time - timestamp > 600:`
`234`	`234`	`raise ValueError("Code has expired")`
`235`		`-`
	`235`	`+`
`236`	`236`	`if stored_code != code:`
`237`	`237`	`raise ValueError("Invalid code")`
`238`	`238`
`@@ -244,35 +244,35 @@ async def verify_email_2fa_code(self, user_id: UUID, code: str) -> bool:`
`244`	`244`	`return True`
`245`	`245`
`246`	`246`	`async def verify_2fa_code(`
`247`		`- self,`
`248`		`- user: User,`
`249`		`- code: str,`
`250`		`- method: Literal["totp", "email", "backup"] = "totp"`
	`247`	`+ self, user: User, code: str, method: Literal["totp", "email", "backup"] = "totp"`
`251`	`248`	`) -> bool:`
`252`	`249`	`"""Verify a 2FA code using the specified method.`
`253`		`-`
	`250`	`+`
`254`	`251`	`Args:`
`255`	`252`	`user: The user entity`
`256`	`253`	`code: The verification code`
`257`	`254`	`method: The verification method ('totp', 'email', or 'backup')`
`258`		`-`
	`255`	`+`
`259`	`256`	`Returns:`
`260`	`257`	`True if verification successful`
`261`	`258`	`"""`
`262`	`259`	`if not user.security:`
`263`	`260`	`raise ValueError("User security not found")`
`264`	`261`
`265`	`262`	`if method == "totp":`
`266`		`- if not user.security.two_factor_secret or not user.security.two_factor_enabled:`
	`263`	`+ if (`
	`264`	`+ not user.security.two_factor_secret`
	`265`	`+ or not user.security.two_factor_enabled`
	`266`	`+ ):`
`267`	`267`	`raise ValueError("TOTP 2FA is not enabled")`
`268`		`-`
	`268`	`+`
`269`	`269`	`totp = pyotp.TOTP(user.security.two_factor_secret)`
`270`	`270`	`return bool(totp.verify(code, valid_window=1))`
`271`	`271`
`272`	`272`	`elif method == "backup":`
`273`	`273`	`if not user.security.two_factor_backup_codes:`
`274`	`274`	`raise ValueError("No backup codes available")`
`275`		`-`
	`275`	`+`
`276`	`276`	`backup_codes = user.security.two_factor_backup_codes.split(",")`
`277`	`277`	`if code in backup_codes:`
`278`	`278`	`# Remove used backup code`
`@@ -289,13 +289,15 @@ async def verify_2fa_code(`
`289`	`289`
`290`	`290`	`return False`
`291`	`291`
`292`		`- async def regenerate_backup_codes(self, user_id: UUID, verify_code: str) -> dict[str, list[str]]:`
	`292`	`+ async def regenerate_backup_codes(`
	`293`	`+ self, user_id: UUID, verify_code: str`
	`294`	`+ ) -> dict[str, list[str]]:`
`293`	`295`	`"""Regenerate backup codes for a user.`
`294`		`-`
	`296`	`+`
`295`	`297`	`Args:`
`296`	`298`	`user_id: The user's ID`
`297`	`299`	`verify_code: Current TOTP code for verification`
`298`		`-`
	`300`	`+`
`299`	`301`	`Returns:`
`300`	`302`	`dict with 'backup_codes' key containing new recovery codes`
`301`	`303`	`"""`

Collapse file

`‎src/modules/user/presentation/dependency.py‎`

Copy file name to clipboardExpand all lines: src/modules/user/presentation/dependency.py

+8-2Lines changed: 8 additions & 2 deletions

Original file line number	Diff line number	Diff line change
`@@ -2,7 +2,8 @@`
`2`	`2`	`from sqlalchemy.ext.asyncio import AsyncSession`
`3`	`3`
`4`	`4`	`from src.core.database.postgres.session import get_db, get_unit_of_work`
`5`		`-from src.core.email.factory import get_email_service`
	`5`	`+from src.core.email.factory import create_email_service`
	`6`	`+from src.core.email.service import EmailService`
`6`	`7`	`from src.core.security.account_lockout import AccountLockoutService`
`7`	`8`	`from src.core.security.audit import AuditService`
`8`	`9`	`from src.core.security.infrastructure.repositories.audit_log_repository import (`
`@@ -53,6 +54,10 @@ def get_user_repository(db: AsyncSession = Depends(get_db)) -> UserRepository:`
`53`	`54`	`return SQLAlchemyUserRepository(db)`
`54`	`55`
`55`	`56`
	`57`	`+def get_email_service() -> EmailService:`
	`58`	`+ return create_email_service()`
	`59`	`+`
	`60`	`+`
`56`	`61`	`def get_refresh_token_repository(`
`57`	`62`	`db: AsyncSession = Depends(get_db),`
`58`	`63`	`) -> RefreshTokenRepository:`
`@@ -128,6 +133,7 @@ def get_logout_handler(`
`128`	`133`
`129`	`134`	`# Two-Factor Authentication Handlers`
`130`	`135`
	`136`	`+`
`131`	`137`	`def get_setup_totp_handler(`
`132`	`138`	`user_repo: UserRepository = Depends(get_user_repository),`
`133`	`139`	`unit_of_work: UnitOfWork = Depends(get_unit_of_work),`
`@@ -152,7 +158,7 @@ def get_disable_totp_handler(`
`152`	`158`	`def get_send_email_2fa_code_handler(`
`153`	`159`	`user_repo: UserRepository = Depends(get_user_repository),`
`154`	`160`	`unit_of_work: UnitOfWork = Depends(get_unit_of_work),`
`155`		`- email_service = Depends(get_email_service),`
	`161`	`+ email_service=Depends(get_email_service),`
`156`	`162`	`) -> SendEmail2FACodeHandler:`
`157`	`163`	`return SendEmail2FACodeHandler(user_repo, unit_of_work, email_service)`
`158`	`164`

0 commit comments

Comments

0 (0)

You are visitor 01402520 — GeezerVison does not store the pages you look up.
Best viewed with Netscape Navigator 4.7 or Internet Explorer 6 at 800×600. No cookies required.