From 3ad7a37f95f5f0a96efeea68b616a9e1e3ff2876 Mon Sep 17 00:00:00 2001
From: Charles Hendrie <chad.hendrie@thomsonreuters.com>
Date: Sat, 8 Oct 2016 19:04:54 -0500
Subject: [PATCH] Fix import certificate private key encoding (#434)

When importing a certificate, the private key is passed to the
import/upload process from the UI as a str object. In Python3 this
raises two issues when processing the private key - the private key
validation fails and database insert of the certificate fails.

The fix in both cases is to correctly encode the private key as a bytes
object.
---
 lemur/certificates/service.py    |  5 +++++
 lemur/common/validators.py       |  8 ++++++--
 lemur/tests/test_certificates.py |  7 +++++++
 lemur/tests/test_validators.py   | 29 +++++++++++++++++++++++++++++
 4 files changed, 47 insertions(+), 2 deletions(-)
 create mode 100644 lemur/tests/test_validators.py

diff --git a/lemur/certificates/service.py b/lemur/certificates/service.py
index cd5786c7..a81aa902 100644
--- a/lemur/certificates/service.py
+++ b/lemur/certificates/service.py
@@ -189,6 +189,11 @@ def upload(**kwargs):
     else:
         kwargs['roles'] = roles
 
+    if kwargs.get('private_key'):
+        private_key = kwargs['private_key']
+        if not isinstance(private_key, bytes):
+            kwargs['private_key'] = private_key.encode('utf-8')
+
     cert = Certificate(**kwargs)
 
     cert = database.create(cert)
diff --git a/lemur/common/validators.py b/lemur/common/validators.py
index 61486dff..7c216402 100644
--- a/lemur/common/validators.py
+++ b/lemur/common/validators.py
@@ -1,5 +1,6 @@
 
-import arrow, re
+import arrow
+import re
 from flask import current_app
 from marshmallow.exceptions import ValidationError
 
@@ -33,7 +34,10 @@ def private_key(key):
     :return: :raise ValueError:
     """
     try:
-        serialization.load_pem_private_key(bytes(key), None, backend=default_backend())
+        if isinstance(key, bytes):
+            serialization.load_pem_private_key(key, None, backend=default_backend())
+        else:
+            serialization.load_pem_private_key(key.encode('utf-8'), None, backend=default_backend())
     except Exception:
         raise ValidationError('Private key presented is not valid.')
 
diff --git a/lemur/tests/test_certificates.py b/lemur/tests/test_certificates.py
index fcf6f887..87e05282 100644
--- a/lemur/tests/test_certificates.py
+++ b/lemur/tests/test_certificates.py
@@ -364,6 +364,13 @@ def test_upload(logged_in_user):
     assert 'ACustomName' in cert.name
 
 
+# verify upload with a private key as a str
+def test_upload_private_key_str(logged_in_user):
+    from lemur.certificates.service import upload
+    cert = upload(body=INTERNAL_VALID_LONG_STR, chain=INTERNAL_VALID_SAN_STR, private_key=PRIVATE_KEY_STR.decode('utf-8'), owner='joe@example.com', name='ACustomName')
+    assert cert
+
+
 @pytest.mark.parametrize("token,status", [
     (VALID_USER_HEADER_TOKEN, 200),
     (VALID_ADMIN_HEADER_TOKEN, 200),
diff --git a/lemur/tests/test_validators.py b/lemur/tests/test_validators.py
new file mode 100644
index 00000000..3dbdaa8f
--- /dev/null
+++ b/lemur/tests/test_validators.py
@@ -0,0 +1,29 @@
+from marshmallow.exceptions import ValidationError
+from .vectors import PRIVATE_KEY_STR
+
+
+def test_private_key():
+    from lemur.common.validators import private_key
+    try:
+        private_key(PRIVATE_KEY_STR)
+        assert True
+    except ValidationError:
+        assert False, "failed to validate private key as a bytes object"
+
+
+def test_private_key_str_object():
+    from lemur.common.validators import private_key
+    try:
+        private_key(PRIVATE_KEY_STR.decode('utf-8'))
+        assert True
+    except ValidationError:
+        assert False, "failed to validate private key as a str object"
+
+
+def test_private_key_invalid():
+    from lemur.common.validators import private_key
+    try:
+        private_key('invalid_private_key')
+        assert False, "invalid private key should have raised an exception"
+    except ValidationError:
+        assert True