Coverage for agentos/tools/totp.py: 0%

1""" 

2TOTP — Time-based One-Time Password (RFC 6238). 

3 

4Supports: 

5 - TOTP generation (SHA1/SHA256/SHA512) 

6 - Configurable digits (6/8) and period (30s default) 

7 - Key URI generation for QR codes (otpauth://) 

8 - Verification with drift tolerance 

9 - HOTP (counter-based) support (RFC 4226) 

10""" 

11 

12from __future__ import annotations 

13 

14import base64 

15import hashlib 

16import hmac 

17import struct 

18import time 

19from urllib.parse import quote, urlencode 

20 

21# ============================================================================ 

22# TOTP / HOTP 

23# ============================================================================ 

24 

25 

26class TOTP: 

27 """Time-based One-Time Password generator (RFC 6238). 

28 

29 Usage: 

30 totp = TOTP(secret="JBSWY3DPEHPK3PXP") 

31 code = totp.now() # 6-digit code 

32 ok = totp.verify(code) # True/False 

33 uri = totp.to_uri("user@example.com", "MyApp") # QR code URI 

34 """ 

35 

36 def __init__( 

37 self, 

38 secret: str, 

39 digits: int = 6, 

40 period: int = 30, 

41 algorithm: str = "SHA1", 

42 ): 

43 self._secret = secret.upper().replace(" ", "") 

44 self._digits = digits 

45 self._period = period 

46 algorithms = {"SHA1": hashlib.sha1, "SHA256": hashlib.sha256, "SHA512": hashlib.sha512} 

47 if algorithm not in algorithms: 

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

49 self._hash_func = algorithms[algorithm] 

50 self._algorithm = algorithm 

51 

52 @property 

53 def secret(self) -> str: 

54 return self._secret 

55 

56 @property 

57 def digits(self) -> int: 

58 return self._digits 

59 

60 @property 

61 def period(self) -> int: 

62 return self._period 

63 

64 @property 

65 def algorithm(self) -> str: 

66 return self._algorithm 

67 

68 # ---------- Generation ---------- 

69 

70 def now(self) -> str: 

71 """Generate the current TOTP code.""" 

72 return self.at(int(time.time())) 

73 

74 def at(self, timestamp: int) -> str: 

75 """Generate TOTP code for a specific Unix timestamp.""" 

76 counter = timestamp // self._period 

77 return self._generate(counter) 

78 

79 # ---------- Verification ---------- 

80 

81 def verify( 

82 self, 

83 code: str, 

84 drift: int = 1, 

85 timestamp: int | None = None, 

86 ) -> bool: 

87 """Verify a TOTP code with optional drift tolerance. 

88 

89 Args: 

90 code: The code to verify 

91 drift: Number of periods before/after to check (default 1 = +/-30s) 

92 timestamp: Reference timestamp, defaults to now 

93 """ 

94 ts = timestamp or int(time.time()) 

95 for offset in range(-drift, drift + 1): 

96 if self.at(ts + offset * self._period) == code: 

97 return True 

98 return False 

99 

100 # ---------- URI ---------- 

101 

102 def to_uri(self, account: str, issuer: str | None = None) -> str: 

103 """Generate otpauth:// URI for QR code. 

104 

105 Args: 

106 account: User account (e.g., email) 

107 issuer: Service name 

108 """ 

109 label = account 

110 if issuer: 

111 label = f"{issuer}:{account}" 

112 

113 params = { 

114 "secret": self._secret, 

115 "digits": str(self._digits), 

116 "period": str(self._period), 

117 "algorithm": self._algorithm, 

118 } 

119 if issuer: 

120 params["issuer"] = issuer 

121 

122 query = urlencode(params) 

123 return f"otpauth://totp/{quote(label)}?{query}" 

124 

125 # ---------- Internal ---------- 

126 

127 def _generate(self, counter: int) -> str: 

128 """Generate HOTP code for a given counter.""" 

129 key = base64.b64decode(self._pad_base64(self._secret)) 

130 msg = struct.pack(">Q", counter) 

131 h = hmac.new(key, msg, self._hash_func).digest() 

132 offset = h[-1] & 0x0F 

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

134 mod = 10**self._digits 

135 return str(binary % mod).zfill(self._digits) 

136 

137 @staticmethod 

138 def _pad_base64(s: str) -> str: 

139 """Pad base32 string for base64 decoding (base32 → base64).""" 

140 # Convert base32 to bytes, then encode as base64 

141 # Standard base32 alphabet: A-Z 2-7, padding = 

142 missing_padding = len(s) % 8 

143 if missing_padding: 

144 s += "=" * (8 - missing_padding) 

145 raw = base64.b32decode(s) 

146 return base64.b64encode(raw).decode("ascii") 

147 

148 @classmethod 

149 def generate_secret(cls, length: int = 32) -> str: 

150 """Generate a random base32 secret.""" 

151 import secrets 

152 

153 raw = secrets.token_bytes(length) 

154 return base64.b32encode(raw).decode("ascii").rstrip("=") 

155 

156 

157class HOTP(TOTP): 

158 """HMAC-based One-Time Password (RFC 4226). 

159 

160 Usage: 

161 hotp = HOTP(secret="JBSWY3DPEHPK3PXP") 

162 code = hotp.at(0) # Generate for counter 0 

163 ok = hotp.verify(code, counter=0) 

164 """ 

165 

166 def __init__( 

167 self, 

168 secret: str, 

169 digits: int = 6, 

170 algorithm: str = "SHA1", 

171 ): 

172 super().__init__(secret=secret, digits=digits, period=1, algorithm=algorithm) 

173 

174 def at(self, counter: int) -> str: 

175 return self._generate(counter) 

176 

177 def verify( 

178 self, 

179 code: str, 

180 counter: int, 

181 look_ahead: int = 10, 

182 ) -> tuple[bool, int | None]: 

183 """Verify HOTP code, returns (is_valid, matched_counter).""" 

184 for c in range(counter, counter + look_ahead + 1): 

185 if self.at(c) == code: 

186 return True, c 

187 return False, None