diff options
Diffstat (limited to 'accounts')
-rw-r--r-- | accounts/forms.py | 3 | ||||
-rw-r--r-- | accounts/utils/sessions.py | 27 |
2 files changed, 17 insertions, 13 deletions
diff --git a/accounts/forms.py b/accounts/forms.py index a85f382..b59ab64 100644 --- a/accounts/forms.py +++ b/accounts/forms.py @@ -1,11 +1,12 @@ # -*- coding: utf-8 -*- import re -from flask import Markup, url_for +from flask import url_for from flask_wtf import FlaskForm as Form from flask_login import current_user from wtforms import StringField, PasswordField, ValidationError, \ BooleanField, validators from wtforms.form import FormMeta +from markupsafe import Markup from .utils import NotRegexp from accounts.app import accounts_app diff --git a/accounts/utils/sessions.py b/accounts/utils/sessions.py index e156f0c..007c928 100644 --- a/accounts/utils/sessions.py +++ b/accounts/utils/sessions.py @@ -3,20 +3,21 @@ from Crypto import Random from Crypto.Cipher import AES -from flask import Flask +from flask import Flask, Request from flask.sessions import TaggedJSONSerializer, SecureCookieSessionInterface from itsdangerous import BadPayload +from base64 import b64encode, b64decode from accounts.app import accounts_app from typing import cast -def _pad(value, block_size): +def _pad(value: str, block_size: int) -> bytes: padding = block_size - len(value) % block_size return (value + (padding * chr(padding))).encode("UTF-8") -def _unpad(value): +def _unpad(value: str) -> str: pad_length = ord(value[len(value)-1:]) return value[:-pad_length] @@ -27,12 +28,12 @@ class EncryptedSerializer(TaggedJSONSerializer): super(EncryptedSerializer, self).__init__() self.block_size = AES.block_size - def _cipher(self, iv): + def _cipher(self, iv: bytes): key = accounts_app.config['SESSION_ENCRYPTION_KEY'] assert len(key) == 32 return AES.new(key, AES.MODE_CBC, iv) - def dumps(self, value): + def dumps(self, value: str) -> str: """ Encrypt the serialized values with `config.SESSION_ENCRYPTION_KEY`. The key must be 32 bytes long. @@ -40,23 +41,25 @@ class EncryptedSerializer(TaggedJSONSerializer): serialized_value = super(EncryptedSerializer, self).dumps(value) raw = _pad(serialized_value, self.block_size) - iv = Random.new().read(self.block_size) - return iv + self._cipher(iv).encrypt(raw) + iv: bytes = Random.new().read(self.block_size) + return b64encode(iv + self._cipher(iv).encrypt(raw)).decode("UTF-8") - def loads(self, value): + def loads(self, value: str): """ Decrypt the given serialized session data with `config.SESSION_ENCRYPTION_KEY`. """ - iv = value[:self.block_size] - raw = self._cipher(iv).decrypt(value[AES.block_size:]) - return super(EncryptedSerializer, self).loads(_unpad(raw)) + decoded = b64decode(value.encode("UTF-8")) + iv = decoded[:self.block_size] + raw = self._cipher(iv).decrypt(decoded[AES.block_size:]) + return super(EncryptedSerializer, self) \ + .loads(_unpad(raw)) class EncryptedSessionInterface(SecureCookieSessionInterface): serializer = EncryptedSerializer() - def open_session(self, app: Flask, request): + def open_session(self, app: Flask, request: Request): session = None try: parent = super(EncryptedSessionInterface, self) |