1
1
"""
2
2
Shared functionality to generate two factor authentication codes
3
3
"""
4
- import base64
5
4
import json
6
5
import os
7
- import otpauth
6
+ from inspect import signature
7
+
8
+ import pyotp
9
+ import qrcode
8
10
from PIL import Image
9
11
from pyzbar .pyzbar import decode
10
- import qrcode
11
-
12
12
13
13
# default configuration file
14
14
config_file = "~/.pyauthenticator"
@@ -75,28 +75,6 @@ def get_otpauth_dict(otpauth_str):
75
75
}
76
76
77
77
78
- def add_padding (main_str , padding_str , padding_length , inverse_padding = False ):
79
- """
80
- Add padding to a string either in the beginning or at the end
81
-
82
- Args:
83
- main_str (str): string to add padding to
84
- padding_str (str): padding character as string
85
- padding_length (int): the length of the final string should be a multiple of the padding length
86
- inverse_padding (bool): add padding in the beginning rather than the end
87
-
88
- Returns:
89
- str: resulting string with padding
90
- """
91
- missing_padding = len (main_str ) % padding_length
92
- if missing_padding :
93
- if inverse_padding :
94
- main_str = padding_str * (padding_length - missing_padding ) + main_str
95
- else :
96
- main_str += padding_str * (padding_length - missing_padding )
97
- return main_str
98
-
99
-
100
78
def check_if_key_in_config (key , config_dict ):
101
79
"""
102
80
Check if a given key is included in a dictionary, raise an ValueError if it is not.
@@ -122,37 +100,25 @@ def get_two_factor_code(key, config_dict):
122
100
"""
123
101
check_if_key_in_config (key = key , config_dict = config_dict )
124
102
decode_dict_internal = get_otpauth_dict (otpauth_str = config_dict [key ])
103
+ funct_sig = signature (pyotp .TOTP )
104
+ if "digits" in decode_dict_internal .keys ():
105
+ digits = int (decode_dict_internal ["digits" ])
106
+ else :
107
+ digits = funct_sig .parameters ["digits" ].default
125
108
if "period" in decode_dict_internal .keys ():
126
- totp = otpauth .TOTP (
127
- secret = base64 .b32decode (
128
- add_padding (
129
- main_str = decode_dict_internal ["secret" ],
130
- padding_str = "=" ,
131
- padding_length = 8 ,
132
- inverse_padding = False ,
133
- ),
134
- True ,
135
- ),
136
- period = int (decode_dict_internal ["period" ]),
137
- )
109
+ interval = int (decode_dict_internal ["period" ])
110
+ else :
111
+ interval = funct_sig .parameters ["interval" ].default
112
+ if "issuer" in decode_dict_internal .keys ():
113
+ issuer = decode_dict_internal ["issuer" ]
138
114
else :
139
- totp = otpauth .TOTP (
140
- secret = base64 .b32decode (
141
- add_padding (
142
- main_str = decode_dict_internal ["secret" ],
143
- padding_str = "=" ,
144
- padding_length = 8 ,
145
- inverse_padding = False ,
146
- ),
147
- True ,
148
- ),
149
- )
150
- return add_padding (
151
- main_str = str (totp .string_code (totp .generate ())),
152
- padding_str = "0" ,
153
- padding_length = 6 ,
154
- inverse_padding = True ,
155
- )
115
+ issuer = funct_sig .parameters ["issuer" ].default
116
+ return pyotp .TOTP (
117
+ s = decode_dict_internal ["secret" ],
118
+ digits = digits ,
119
+ issuer = issuer ,
120
+ interval = interval ,
121
+ ).now ()
156
122
157
123
158
124
def add_service (
0 commit comments