From 2c7cb4a370bb847b43b98fd09ef358cd73af6daa Mon Sep 17 00:00:00 2001
From: Ryan Petrello <rpetrell@redhat.com>
Date: Fri, 10 Feb 2017 14:30:42 -0500
Subject: [PATCH] add utf-8 support to utils.common.encrypt_field/decrypt_field

---
 .../tests/unit/utils/common/test_common.py    | 19 ++++++++++++---
 awx/main/utils/common.py                      | 23 ++++++++++++++++---
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/awx/main/tests/unit/utils/common/test_common.py b/awx/main/tests/unit/utils/common/test_common.py
index a152d69c12..a48bbe64b3 100644
--- a/awx/main/tests/unit/utils/common/test_common.py
+++ b/awx/main/tests/unit/utils/common/test_common.py
@@ -1,24 +1,37 @@
+# -*- coding: utf-8 -*-
+
+# Copyright (c) 2017 Ansible, Inc.
+# All Rights Reserved.
+
 from awx.conf.models import Setting
 from awx.main.utils import common
 
 
 def test_encrypt_field():
     field = Setting(pk=123, value='ANSIBLE')
-    encrypted = common.encrypt_field(field, 'value')
+    encrypted = field.value = common.encrypt_field(field, 'value')
     assert encrypted == '$encrypted$AES$Ey83gcmMuBBT1OEq2lepnw=='
     assert common.decrypt_field(field, 'value') == 'ANSIBLE'
 
 
 def test_encrypt_field_without_pk():
     field = Setting(value='ANSIBLE')
-    encrypted = common.encrypt_field(field, 'value')
+    encrypted = field.value = common.encrypt_field(field, 'value')
     assert encrypted == '$encrypted$AES$8uIzEoGyY6QJwoTWbMFGhw=='
     assert common.decrypt_field(field, 'value') == 'ANSIBLE'
 
 
+def test_encrypt_field_with_unicode_string():
+    value = u'Iñtërnâtiônàlizætiøn'
+    field = Setting(value=value)
+    encrypted = field.value = common.encrypt_field(field, 'value')
+    assert encrypted == '$encrypted$UTF8$AES$AESQbqOefpYcLC7x8yZ2aWG4FlXlS66JgavLbDp/DSM='
+    assert common.decrypt_field(field, 'value') == value
+
+
 def test_encrypt_subfield():
     field = Setting(value={'name': 'ANSIBLE'})
-    encrypted = common.encrypt_field(field, 'value', subfield='name')
+    encrypted = field.value = common.encrypt_field(field, 'value', subfield='name')
     assert encrypted == '$encrypted$AES$8uIzEoGyY6QJwoTWbMFGhw=='
     assert common.decrypt_field(field, 'value', subfield='name') == 'ANSIBLE'
 
diff --git a/awx/main/utils/common.py b/awx/main/utils/common.py
index 5842b78db9..e49f4d0131 100644
--- a/awx/main/utils/common.py
+++ b/awx/main/utils/common.py
@@ -21,6 +21,8 @@ import tempfile
 # Decorator
 from decorator import decorator
 
+import six
+
 # Django
 from django.utils.translation import ugettext_lazy as _
 from django.db.models import ManyToManyField
@@ -190,6 +192,7 @@ def encrypt_field(instance, field_name, ask=False, subfield=None):
         value = value[subfield]
     if not value or value.startswith('$encrypted$') or (ask and value == 'ASK'):
         return value
+    utf8 = type(value) == six.text_type
     value = smart_str(value)
     key = get_encryption_key(field_name, getattr(instance, 'pk', None))
     cipher = AES.new(key, AES.MODE_ECB)
@@ -197,17 +200,31 @@ def encrypt_field(instance, field_name, ask=False, subfield=None):
         value += '\x00'
     encrypted = cipher.encrypt(value)
     b64data = base64.b64encode(encrypted)
-    return '$encrypted$%s$%s' % ('AES', b64data)
+    tokens = ['$encrypted', 'AES', b64data]
+    if utf8:
+        # If the value to encrypt is utf-8, we need to add a marker so we
+        # know to decode the data when it's decrypted later
+        tokens.insert(1, 'UTF8')
+    return '$'.join(tokens)
 
 
 def decrypt_value(encryption_key, value):
-    algo, b64data = value[len('$encrypted$'):].split('$', 1)
+    raw_data = value[len('$encrypted$'):]
+    # If the encrypted string contains a UTF8 marker, discard it
+    utf8 = raw_data.startswith('UTF8$')
+    if utf8:
+        raw_data = raw_data[len('UTF8$'):]
+    algo, b64data = raw_data.split('$', 1)
     if algo != 'AES':
         raise ValueError('unsupported algorithm: %s' % algo)
     encrypted = base64.b64decode(b64data)
     cipher = AES.new(encryption_key, AES.MODE_ECB)
     value = cipher.decrypt(encrypted)
-    return value.rstrip('\x00')
+    value = value.rstrip('\x00')
+    # If the encrypted string contained a UTF8 marker, decode the data
+    if utf8:
+        value = value.decode('utf-8')
+    return value
 
 
 def decrypt_field(instance, field_name, subfield=None):