Coverage for src / kdbxtool / security / totp.py: 92%

1"""TOTP (Time-based One-Time Password) implementation per RFC 6238. 

2 

3This module provides TOTP code generation from otpauth:// URIs stored in 

4KeePass entry otp fields. 

5""" 

6 

7from __future__ import annotations 

8 

9import base64 

10import hashlib 

11import hmac 

12import struct 

13import time 

14from dataclasses import dataclass 

15from datetime import UTC, datetime 

16from typing import Literal 

17from urllib.parse import parse_qs, unquote, urlparse 

18 

19 

20@dataclass 

21class TotpCode: 

22 """A generated TOTP code with expiration info. 

23 

24 Attributes: 

25 code: The TOTP code as a zero-padded string (e.g., "123456") 

26 period: Time step in seconds (typically 30) 

27 generated_at: Unix timestamp when code was generated 

28 """ 

29 

30 code: str 

31 period: int 

32 generated_at: float 

33 

34 @property 

35 def remaining(self) -> int: 

36 """Seconds remaining until this code expires. 

37 

38 Note: This is calculated fresh each call based on current time. 

39 """ 

40 now = time.time() 

41 elapsed = now - (int(self.generated_at) // self.period * self.period) 

42 return max(0, self.period - int(elapsed)) 

43 

44 @property 

45 def expires_at(self) -> datetime: 

46 """Datetime when this code expires.""" 

47 period_start = int(self.generated_at) // self.period * self.period 

48 return datetime.fromtimestamp(period_start + self.period, tz=UTC) 

49 

50 @property 

51 def is_expired(self) -> bool: 

52 """Whether this code has expired.""" 

53 return self.remaining <= 0 

54 

55 def __str__(self) -> str: 

56 return self.code 

57 

58 

59@dataclass 

60class TotpConfig: 

61 """TOTP configuration parsed from an otpauth:// URI. 

62 

63 Attributes: 

64 secret: Base32-encoded secret key (decoded to bytes internally) 

65 digits: Number of digits in the code (default: 6) 

66 period: Time step in seconds (default: 30) 

67 algorithm: Hash algorithm (SHA1, SHA256, or SHA512) 

68 issuer: Optional issuer name 

69 account: Optional account name/label 

70 """ 

71 

72 secret: bytes 

73 digits: int = 6 

74 period: int = 30 

75 algorithm: Literal["SHA1", "SHA256", "SHA512"] = "SHA1" 

76 issuer: str | None = None 

77 account: str | None = None 

78 

79 

80def parse_otpauth_uri(uri: str) -> TotpConfig: 

81 """Parse an otpauth:// URI into a TotpConfig. 

82 

83 Supports the standard otpauth:// URI format: 

84 otpauth://totp/LABEL?secret=BASE32SECRET&issuer=ISSUER&... 

85 

86 Args: 

87 uri: The otpauth:// URI string 

88 

89 Returns: 

90 TotpConfig with parsed parameters 

91 

92 Raises: 

93 ValueError: If the URI is invalid or missing required parameters 

94 """ 

95 parsed = urlparse(uri) 

96 

97 if parsed.scheme != "otpauth": 

98 raise ValueError(f"Invalid scheme: expected 'otpauth', got '{parsed.scheme}'") 

99 

100 otp_type = parsed.netloc.lower() 

101 if otp_type not in ("totp", "hotp"): 

102 raise ValueError(f"Unsupported OTP type: {otp_type}") 

103 

104 if otp_type == "hotp": 

105 raise ValueError("HOTP is not supported, only TOTP") 

106 

107 # Parse query parameters 

108 params = parse_qs(parsed.query) 

109 

110 # Secret is required 

111 if "secret" not in params: 

112 raise ValueError("Missing required 'secret' parameter") 

113 

114 secret_b32 = params["secret"][0].upper() 

115 # Add padding if needed (base32 requires padding to multiple of 8) 

116 padding = (8 - len(secret_b32) % 8) % 8 

117 secret_b32 += "=" * padding 

118 

119 try: 

120 secret = base64.b32decode(secret_b32) 

121 except Exception as e: 

122 raise ValueError(f"Invalid base32 secret: {e}") from e 

123 

124 # Parse optional parameters with defaults 

125 digits = int(params.get("digits", ["6"])[0]) 

126 if digits not in (6, 7, 8): 

127 raise ValueError(f"Invalid digits: {digits} (must be 6, 7, or 8)") 

128 

129 period = int(params.get("period", ["30"])[0]) 

130 if period <= 0: 

131 raise ValueError(f"Invalid period: {period} (must be positive)") 

132 

133 algorithm = params.get("algorithm", ["SHA1"])[0].upper() 

134 if algorithm not in ("SHA1", "SHA256", "SHA512"): 

135 raise ValueError(f"Unsupported algorithm: {algorithm}") 

136 

137 issuer = params.get("issuer", [None])[0] 

138 

139 # Parse label (path component) for account name 

140 label = unquote(parsed.path.lstrip("/")) 

141 account = None 

142 if label: 

143 # Label format: "Issuer:Account" or just "Account" 

144 if ":" in label: 

145 _, account = label.split(":", 1) 

146 else: 

147 account = label 

148 

149 return TotpConfig( 

150 secret=secret, 

151 digits=digits, 

152 period=period, 

153 algorithm=algorithm, # type: ignore[arg-type] 

154 issuer=issuer, 

155 account=account, 

156 ) 

157 

158 

159def parse_keepassxc_legacy(seed: str, settings: str | None = None) -> TotpConfig: 

160 """Parse KeePassXC legacy TOTP fields. 

161 

162 KeePassXC historically stored TOTP in separate custom fields: 

163 - "TOTP Seed": Base32 secret 

164 - "TOTP Settings": "period;digits" (e.g., "30;6") 

165 

166 Args: 

167 seed: The TOTP seed (base32 encoded secret) 

168 settings: Optional settings string in "period;digits" format 

169 

170 Returns: 

171 TotpConfig with parsed parameters 

172 """ 

173 # Clean up seed 

174 secret_b32 = seed.strip().upper().replace(" ", "") 

175 padding = (8 - len(secret_b32) % 8) % 8 

176 secret_b32 += "=" * padding 

177 

178 try: 

179 secret = base64.b32decode(secret_b32) 

180 except Exception as e: 

181 raise ValueError(f"Invalid base32 seed: {e}") from e 

182 

183 period = 30 

184 digits = 6 

185 

186 if settings: 

187 parts = settings.split(";") 

188 if len(parts) >= 1 and parts[0]: 

189 period = int(parts[0]) 

190 if len(parts) >= 2 and parts[1]: 

191 digits = int(parts[1]) 

192 

193 return TotpConfig(secret=secret, digits=digits, period=period) 

194 

195 

196def generate_totp(config: TotpConfig, timestamp: float | None = None) -> TotpCode: 

197 """Generate a TOTP code. 

198 

199 Args: 

200 config: TOTP configuration 

201 timestamp: Unix timestamp (defaults to current time) 

202 

203 Returns: 

204 TotpCode with code string and expiration info 

205 """ 

206 if timestamp is None: 

207 timestamp = time.time() 

208 

209 # Calculate time counter 

210 counter = int(timestamp) // config.period 

211 

212 # Select hash algorithm 

213 if config.algorithm == "SHA1": 

214 digest = hashlib.sha1 

215 elif config.algorithm == "SHA256": 

216 digest = hashlib.sha256 

217 else: # SHA512 

218 digest = hashlib.sha512 

219 

220 # Compute HMAC 

221 counter_bytes = struct.pack(">Q", counter) 

222 mac = hmac.new(config.secret, counter_bytes, digest).digest() 

223 

224 # Dynamic truncation (RFC 4226) 

225 offset = mac[-1] & 0x0F 

226 binary = struct.unpack(">I", mac[offset : offset + 4])[0] & 0x7FFFFFFF 

227 

228 # Generate code with specified digits 

229 code_int = binary % (10**config.digits) 

230 code = str(code_int).zfill(config.digits) 

231 

232 return TotpCode(code=code, period=config.period, generated_at=timestamp)