-
Notifications
You must be signed in to change notification settings - Fork 0
/
flask_totp.py
66 lines (53 loc) · 2.29 KB
/
flask_totp.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
from flask import current_app, jsonify, request
from functools import wraps
import pyotp
__version__ = '1.0.0'
__extension_key__ = 'FLASK_TOTP'
__secret_key__ = 'FLASK_TOTP_SECRET'
class FlaskTOTP(object):
"""
Wrapper around PyOTP library for handling time-based OTP counters in a Flask environment.
"""
def __init__(self, app=None):
"""
:param app: optional Flask application object (defaults to None)
"""
if app is not None:
self.init_app(app)
def init_app(self, app):
"""
:param app: the Flask application object
:raises RuntimeError: if seed 'FLASK_TOTP_SECRET' is not provided in configuration
"""
secret = app.config.get(__secret_key__)
if secret is None:
raise RuntimeError("{} not defined in configuration".format(__secret_key__))
self.totp = pyotp.TOTP(secret)
if not hasattr(app, 'extensions'):
app.extensions = {}
app.extensions[__extension_key__] = self
def verify(self, token, for_time=None, valid_window=0):
"""
Verifies the OTP passed in against the current time OTP.
:param otp: the OTP to check against
:param for_time: Time to check OTP at (defaults to now)
:param valid_window: extends the validity to this many counter ticks before and after the current one
:returns: True if verification succeeded, False otherwise
"""
return self.totp.verify(token, for_time, valid_window)
def totp_required(f):
"""
Decorates route to verify the OTP passed with authorization header against the current time OTP before chain proceeds.
Request Header -> Authorization: TOTP 528941
:returns: following view function if verification succeeded, else jsonify with message for failure
"""
@wraps(f)
def decorated_function(*args, **kwargs):
auth_header = request.headers.get('Authorization')
if auth_header is not None:
scheme, token = auth_header.split()
flask_totp = current_app.extensions.get(__extension_key__)
if flask_totp and scheme == 'TOTP' and flask_totp.verify(token) == True:
return f(*args, **kwargs)
return jsonify(message = "TOTP verification failed.")
return decorated_function