Coverage for gws-app/gws/lib/otp/__init__.py: 96%

1"""Generate HOTP and TOTP tokens. 

2 

3References: 

4 https://datatracker.ietf.org/doc/html/rfc4226 

5 https://datatracker.ietf.org/doc/html/rfc6238 

6""" 

7 

8from typing import Optional, cast 

9 

10import base64 

11import hashlib 

12import hmac 

13import random 

14 

15import gws 

16import gws.lib.net 

17 

18 

19class Options(gws.Data): 

20 start: int 

21 step: int 

22 length: int 

23 tolerance: int 

24 algo: str 

25 

26 

27DEFAULTS = Options( 

28 start=0, 

29 step=30, 

30 length=6, 

31 tolerance=1, 

32 algo='sha1', 

33) 

34 

35 

36def new_hotp(secret: str | bytes, counter: int, options: Optional[Options] = None) -> str: 

37 """Generate a new HOTP value as per rfc4226 section 5.3.""" 

38 

39 options = cast(Options, gws.u.merge(DEFAULTS, options)) 

40 return _raw_otp(_to_bytes(secret), counter, options) 

41 

42 

43def new_totp(secret: str | bytes, timestamp: int, options: Optional[Options] = None) -> str: 

44 """Generate a new TOTP value as per rfc6238 section 4.2.""" 

45 

46 options = cast(Options, gws.u.merge(DEFAULTS, options)) 

47 counter = (timestamp - options.start) // options.step 

48 return _raw_otp(_to_bytes(secret), counter, options) 

49 

50 

51def check_totp(input: str, secret: str, timestamp: int, options: Optional[Options] = None) -> bool: 

52 """Check if the input TOTP is valid. 

53 

54 Compares the input against several TOTPs within the tolerance window 

55 ``(timestamp-step*tolerance...timestamp+step*tolerance)``. 

56 """ 

57 

58 options = cast(Options, gws.u.merge(DEFAULTS, options)) 

59 

60 if len(input) != options.length: 

61 return False 

62 

63 ok = False 

64 

65 for window in range(-options.tolerance, options.tolerance + 1): 

66 ts = timestamp + options.step * window 

67 counter = (ts - options.start) // options.step 

68 totp = _raw_otp(_to_bytes(secret), counter, options) 

69 if hmac.compare_digest(_to_bytes(input), _to_bytes(totp)): 

70 ok = True 

71 

72 return ok 

73 

74 

75def totp_key_uri( 

76 secret: str | bytes, 

77 issuer_name: str, 

78 account_name: str, 

79 options: Optional[Options] = None 

80) -> str: 

81 return _key_uri('totp', secret, issuer_name, account_name, None, options) 

82 

83 

84def hotp_key_uri( 

85 secret: str | bytes, 

86 issuer_name: str, 

87 account_name: str, 

88 counter: int, 

89 options: Optional[Options] = None 

90) -> str: 

91 return _key_uri('hotp', secret, issuer_name, account_name, counter, options) 

92 

93 

94def _key_uri( 

95 method: str, 

96 secret: str | bytes, 

97 issuer_name: str, 

98 account_name: str, 

99 counter: Optional[int] = None, 

100 options: Optional[Options] = None 

101) -> str: 

102 """Create a key uri for auth apps. 

103 

104 Reference: 

105 https://github.com/google/google-authenticator/wiki/Key-Uri-Format 

106 """ 

107 

108 params: dict = { 

109 'secret': base32_encode(secret), 

110 'issuer': issuer_name, 

111 } 

112 

113 options = cast(Options, gws.u.merge(DEFAULTS, options)) 

114 

115 if options.algo != DEFAULTS.algo: 

116 params['algorithm'] = options.algo 

117 if options.length != DEFAULTS.length: 

118 params['digits'] = options.length 

119 

120 if method == 'hotp': 

121 params['counter'] = counter 

122 elif options.step != DEFAULTS.step: 

123 params['period'] = options.step 

124 

125 return 'otpauth://{}/{}:{}?{}'.format( 

126 method, 

127 gws.lib.net.quote_param(issuer_name), 

128 gws.lib.net.quote_param(account_name), 

129 gws.lib.net.make_qs(params) 

130 ) 

131 

132 

133def base32_decode(s: str) -> bytes: 

134 return base64.b32decode(s) 

135 

136 

137def base32_encode(s: str | bytes) -> str: 

138 return base64.b32encode(_to_bytes(s)).decode('ascii') 

139 

140 

141def random_secret(base32_length: int = 32) -> str: 

142 """Generate a random printable secret that fits into base32_length.""" 

143 

144 if (base32_length & 7) != 0: 

145 raise ValueError('invalid length') 

146 

147 size = (base32_length >> 3) * 5 

148 r = random.SystemRandom() 

149 return ''.join(chr(r.randint(0x21, 0x7f)) for _ in range(size)) 

150 

151 

152## 

153 

154def _raw_otp(key: bytes, counter: int, options: Options) -> str: 

155 # https://www.rfc-editor.org/rfc/rfc4226#section-5.3 

156 # 

157 # Step 1: Generate an HMAC-SHA-1 value 

158 # Let HS = HMAC-SHA-1(K,C) // HS is a 20-byte string 

159 # 

160 # Step 2: Generate a 4-byte string (Dynamic Truncation) 

161 # Let Sbits = DT(HS) // DT, defined below, returns a 31-bit string 

162 # 

163 # Let OffsetBits be the low-order 4 bits of String[19] 

164 # Offset = StToNum(OffsetBits) // 0 <= OffSet <= 15 

165 # Let P = String[OffSet]...String[OffSet+3] 

166 # Return the Last 31 bits of P 

167 # 

168 # Let Snum = StToNum(Sbits) // Convert S to a number in 0...2^{31}-1 

169 # 

170 # Step 3: Compute an HOTP value 

171 # Return D = Snum mod 10^Digit // D is a number in the range 0...10^{Digit}-1 

172 

173 c = counter.to_bytes(8, byteorder='big') 

174 

175 digestmod = getattr(hashlib, options.algo.lower()) 

176 hs = hmac.new(key, c, digestmod).digest() 

177 

178 offset = hs[-1] & 0xf 

179 p = hs[offset:offset + 4] 

180 snum = int.from_bytes(p, byteorder='big', signed=False) & 0x7fffffff 

181 

182 d = snum % (10 ** options.length) 

183 

184 return f'{d:0{options.length}d}' 

185 

186 

187def _to_bytes(s): 

188 return s.encode('utf8') if isinstance(s, str) else s 

189 

190 

191def _option(options, key, default): 

192 if not options: 

193 return default 

194 return getattr(options, key, default)