forked from lingthio/Flask-User
-
Notifications
You must be signed in to change notification settings - Fork 0
/
tokens.py
80 lines (71 loc) · 3.18 KB
/
tokens.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
67
68
69
70
71
72
73
74
75
76
77
78
79
80
""" This file contains functions to generate and verify tokens for Flask-User.
Tokens contain an encoded user ID and a signature. The signature is managed by the itsdangerous module.
:copyright: (c) 2013 by Ling Thio
:author: Ling Thio ([email protected])
:license: Simplified BSD License, see LICENSE.txt for more details."""
import base64
from Crypto.Cipher import AES
from itsdangerous import BadSignature, SignatureExpired, TimestampSigner
class TokenManager(object):
def setup(self, secret):
""" Create a cypher to encrypt IDs and a signer to sign tokens."""
# Create cypher to encrypt IDs
# and ensure >=16 characters
precursor = b'0123456789abcdef'
if isinstance(secret, bytes):
key = secret + precursor
else:
key = secret.encode("utf-8") + precursor
self.cipher = AES.new(key[0:16], AES.MODE_ECB)
# Create signer to sign tokens
self.signer = TimestampSigner(secret)
def encrypt_id(self, id):
""" Encrypts integer ID to url-safe base64 string."""
# 16 byte integer
str1 = '%016d' % id
# encrypted data
str2 = self.cipher.encrypt(str1.encode())
# URL safe base64 string with '=='
str3 = base64.urlsafe_b64encode(str2)
# return base64 string without '=='
return str3[0:-2]
def decrypt_id(self, encrypted_id):
""" Decrypts url-safe base64 string to integer ID"""
# Convert strings and unicode strings to bytes if needed
if hasattr(encrypted_id, 'encode'):
encrypted_id = encrypted_id.encode('ascii', 'ignore')
try:
str3 = encrypted_id + b'==' # --> base64 string with '=='
#print('str3=', str3)
str2 = base64.urlsafe_b64decode(str3) # --> encrypted data
#print('str2=', str2)
str1 = self.cipher.decrypt(str2) # --> 16 byte integer string
#print('str1=', str1)
return int(str1) # --> integer id
except Exception as e: # pragma: no cover
print('!!!Exception in decrypt_id!!!:', e)
return 0
def generate_token(self, id):
""" Return token with id, timestamp and signature"""
# In Python3 we must make sure that bytes are converted to strings.
# Hence the addition of '.decode()'
return self.signer.sign(self.encrypt_id(id)).decode()
def verify_token(self, token, expiration_in_seconds):
""" Verify token and return (is_valid, has_expired, id).
Returns (True, False, id) on success.
Returns (False, True, None) on expired tokens.
Returns (False, False, None) on invalid tokens."""
try:
data = self.signer.unsign(token, max_age=expiration_in_seconds)
is_valid = True
has_expired = False
id = self.decrypt_id(data)
except SignatureExpired:
is_valid = False
has_expired = True
id = None
except BadSignature:
is_valid = False
has_expired = False
id = None
return (is_valid, has_expired, id)