Coverage for src/lexigram/auth/mfa/manager.py: 94%

1"""MFA Manager — consolidated TOTP and backup-code management. 

2 

3Stores MFA state in the user's ``profile['mfa']`` dict so no separate 

4MFA table is required. Profile keys: 

5 

6- ``enabled``: bool — whether TOTP is currently active 

7- ``secret``: base32-encoded TOTP secret 

8- ``backup_codes``: list of SHA-256 hex digests of the one-time codes 

9""" 

10 

11from __future__ import annotations 

12 

13import dataclasses 

14import hashlib 

15from typing import TYPE_CHECKING 

16 

17from lexigram.auth.authn.mfa import ( 

18 generate_backup_codes, 

19 generate_totp_secret, 

20 get_provisioning_uri, 

21 hash_backup_codes, 

22 verify_totp, 

23) 

24from lexigram.di.decorators import inject 

25from lexigram.logging import get_logger 

26 

27if TYPE_CHECKING: 

28 from lexigram.auth.storage.token_store import UserStoreProtocol 

29 

30logger = get_logger(__name__) 

31 

32 

33@inject 

34class MFAManager: 

35 """Manages Multi-Factor Authentication (TOTP + backup codes) for users. 

36 

37 MFA state is stored in the user profile so no dedicated MFA table is 

38 needed. Injected via the DI container; the resolved ``UserStoreProtocol`` is 

39 used for all persistence. 

40 

41 Args: 

42 user_store: The user store used to read and persist user profiles. 

43 """ 

44 

45 def __init__(self, user_store: UserStoreProtocol) -> None: 

46 self.user_store = user_store 

47 

48 async def enable_totp( 

49 self, 

50 user_id: str, 

51 issuer: str = "lexigram", 

52 ) -> tuple[str, str, list[str]]: 

53 """Enable TOTP for a user and return enrollment credentials. 

54 

55 A fresh TOTP secret and a set of backup codes are generated and 

56 stored (only the SHA-256 digests of backup codes are persisted). 

57 

58 Args: 

59 user_id: The user to enable TOTP for. 

60 issuer: The issuer label shown in authenticator apps. 

61 

62 Returns: 

63 A ``(secret, provisioning_uri, plain_backup_codes)`` tuple. 

64 Show the plain backup codes to the user once — they are not 

65 stored in plaintext. 

66 

67 Raises: 

68 ValueError: If the user does not exist. 

69 """ 

70 user = await self.user_store.get_user_by_id(user_id) 

71 if not user: 

72 raise ValueError("User not found") 

73 

74 secret = generate_totp_secret() 

75 account_name = user.name or user.email or user.user_id 

76 provisioning_uri = get_provisioning_uri(secret, account_name, issuer) 

77 backup_codes = generate_backup_codes() 

78 backup_hashes = hash_backup_codes(backup_codes) 

79 

80 profile = dict(user.profile) 

81 profile_mfa = dict(profile.get("mfa") or {}) 

82 profile_mfa.update( 

83 {"enabled": True, "secret": secret, "backup_codes": backup_hashes}, 

84 ) 

85 profile["mfa"] = profile_mfa 

86 updated = dataclasses.replace(user, profile=profile) 

87 await self.user_store.update_user(updated) 

88 

89 return secret, provisioning_uri, backup_codes 

90 

91 async def verify_totp(self, user_id: str, code: str) -> bool: 

92 """Verify a TOTP or backup code for a user. 

93 

94 A backup code is consumed on first use — a second attempt with the 

95 same backup code returns ``False``. 

96 

97 Args: 

98 user_id: The user to verify. 

99 code: A 6-digit TOTP code or a raw backup code string. 

100 

101 Returns: 

102 ``True`` if the code is valid, ``False`` otherwise. 

103 """ 

104 user = await self.user_store.get_user_by_id(user_id) 

105 if not user: 

106 return False 

107 

108 mfa = user.profile.get("mfa") or {} 

109 if not mfa.get("enabled"): 

110 return False 

111 

112 secret = mfa.get("secret") 

113 if secret and verify_totp(secret, code): 

114 return True 

115 

116 # Check backup codes (single-use). 

117 backup_hashes = list(mfa.get("backup_codes") or []) 

118 code_hash = hashlib.sha256(str(code).encode("utf-8")).hexdigest() 

119 if code_hash in backup_hashes: 

120 backup_hashes.remove(code_hash) 

121 profile = dict(user.profile) 

122 profile_mfa = dict(profile.get("mfa") or {}) 

123 profile_mfa["backup_codes"] = backup_hashes 

124 profile["mfa"] = profile_mfa 

125 updated = dataclasses.replace(user, profile=profile) 

126 await self.user_store.update_user(updated) 

127 return True 

128 

129 return False 

130 

131 async def disable_totp(self, user_id: str) -> bool: 

132 """Disable TOTP for a user. 

133 

134 Args: 

135 user_id: The user to disable TOTP for. 

136 

137 Returns: 

138 ``True`` if TOTP was disabled, ``False`` if the user does not exist. 

139 """ 

140 user = await self.user_store.get_user_by_id(user_id) 

141 if not user: 

142 return False 

143 

144 profile = dict(user.profile) 

145 profile_mfa = dict(profile.get("mfa") or {}) 

146 profile_mfa.update({"enabled": False, "secret": None, "backup_codes": []}) 

147 profile["mfa"] = profile_mfa 

148 updated = dataclasses.replace(user, profile=profile) 

149 await self.user_store.update_user(updated) 

150 return True 

151 

152 def __repr__(self) -> str: 

153 """Return a string representation of this MFAManager.""" 

154 return f"MFAManager(store={type(self.user_store).__name__})" 

155 

156 

157__all__ = ["MFAManager"]