From 0a48997fdd32b4027e07361dbd4b2238fbdae62e Mon Sep 17 00:00:00 2001
From: Matt Wright <matt@nobien.net>
Date: Tue, 10 Jun 2014 11:47:35 -0400
Subject: [PATCH] Improve encoding of strings. Addresses #231 and #253

---
 flask_security/utils.py | 18 +++++++++++++++---
 tests/test_misc.py      | 18 +++++++++++++++++-
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/flask_security/utils.py b/flask_security/utils.py
index b9837c42..f519b3d0 100644
--- a/flask_security/utils.py
+++ b/flask_security/utils.py
@@ -99,13 +99,15 @@ def get_hmac(password):
 
     :param password: The password to sign
     """
-    if _security.password_salt is None:
+    salt = _security.password_salt
+
+    if salt is None:
         raise RuntimeError(
             'The configuration value `SECURITY_PASSWORD_SALT` must '
             'not be None when the value of `SECURITY_PASSWORD_HASH` is '
             'set to "%s"' % _security.password_hash)
 
-    h = hmac.new(_security.password_salt.encode('utf-8'), password.encode('utf-8'), hashlib.sha512)
+    h = hmac.new(encode_string(salt), encode_string(password), hashlib.sha512)
     return base64.b64encode(h.digest())
 
 
@@ -149,8 +151,18 @@ def encrypt_password(password):
     return _pwd_context.encrypt(signed)
 
 
+def encode_string(string):
+    """Encodes a string to bytes, if it isn't already.
+
+    :param string: The string to encode"""
+
+    if isinstance(string, text_type):
+        string = string.encode('utf-8')
+    return string
+
+
 def md5(data):
-    return hashlib.md5(data.encode('utf-8')).hexdigest()
+    return hashlib.md5(encode_string(data)).hexdigest()
 
 
 def do_flash(message, category=None):
diff --git a/tests/test_misc.py b/tests/test_misc.py
index 8be2a7d3..9b76457a 100644
--- a/tests/test_misc.py
+++ b/tests/test_misc.py
@@ -12,7 +12,7 @@
 from flask_security.forms import LoginForm, RegisterForm, ConfirmRegisterForm, \
     SendConfirmationForm, PasswordlessLoginForm, ForgotPasswordForm, ResetPasswordForm, \
     ChangePasswordForm, TextField, PasswordField, email_required, email_validator, valid_user_email
-from flask_security.utils import capture_reset_password_requests
+from flask_security.utils import capture_reset_password_requests, md5, string_types
 
 from utils import authenticate, init_app_with_options, populate_data
 
@@ -170,3 +170,19 @@ def test_change_hash_type(app, sqlalchemy_datastore):
 
     response = client.post('/login', data=dict(email='matt@lp.com', password='password'))
     assert response.status_code == 302
+
+
+def test_md5():
+    data = md5(b'hello')
+    assert isinstance(data, string_types)
+    data = md5(u'hellö')
+    assert isinstance(data, string_types)
+
+
+@pytest.mark.settings(password_salt=u'öööööööööööööööööööööööööööööööööö',
+                      password_hash='bcrypt')
+def test_password_unicode_password_salt(client):
+    response = authenticate(client)
+    assert response.status_code == 302
+    response = authenticate(client, follow_redirects=True)
+    assert b'Hello matt@lp.com' in response.data