From af4d9558069dc88b9daf95d6085a62fad9753e47 Mon Sep 17 00:00:00 2001
From: Jakub Warmuz <jakub@warmuz.org>
Date: Sun, 23 Nov 2014 02:37:34 +0100
Subject: [PATCH] Fix JOSE encoding mess

---
 letsencrypt/client/le_util.py      | 53 +++++++++++++++---------------
 letsencrypt/client/le_util_test.py | 26 +++++++--------
 2 files changed, 39 insertions(+), 40 deletions(-)

diff --git a/letsencrypt/client/le_util.py b/letsencrypt/client/le_util.py
index 68a8c405d..31edfd13b 100644
--- a/letsencrypt/client/le_util.py
+++ b/letsencrypt/client/le_util.py
@@ -65,48 +65,47 @@ def unique_file(default_name, mode=0777):
         count += 1
 
 
-def _to_utf8(arg):
-    """Normalize to UTF-8 string."""
-    return arg.encode('utf-8') if isinstance(arg, unicode) else arg
+# https://tools.ietf.org/html/draft-ietf-jose-json-web-signature-37#appendix-C
+#
+# Jose Base64:
+#
+#   - URL-safe Base64
+#
+#   - padding stripped
 
 
-def jose_b64encode(arg):
+def jose_b64encode(data, encoding='utf-8'):
     """JOSE Base64 encode.
 
-    JOSE Base64:
-    - URL-safe Base64
-    - padding stripped
+    :param data: Data to be encoded.
+    :type data: str or unicode
 
-    https://tools.ietf.org/html/draft-ietf-jose-json-web-signature-37#appendix-C
-
-    :param arg: String to be encoded. Unicode input will be encoded
-                to UTF-8 before Base64 encoding.
-    :type arg: str or unicode
+    :param encoding: Name of the encoding to be performed before
+                     Base64 encoding. If not None, then `data`
+                     has to be unicode.
+    :type encoding: str or None
 
     :returns: JOSE Base64 string.
     :rtype: str
 
     """
-    return base64.urlsafe_b64encode(_to_utf8(arg)).rstrip('=')
+    encoded = data if encoding is None else data.encode(encoding)
+    return base64.urlsafe_b64encode(encoded).rstrip('=')
 
 
-def jose_b64decode(arg):
+def jose_b64decode(arg, decoding='utf-8'):
     """JOSE Base64 decode.
 
-    Jose Base64:
-    - URL-safe Base64
-    - padding stripped
+    :param arg: Base64 string to be decoded.
+    :type arg: str
 
-    https://tools.ietf.org/html/draft-ietf-jose-json-web-signature-37#appendix-C
+    :param decoding: Name of the encoding to be performed after
+                     Base64 decoding.
+    :type decoding: str or None
 
-    :param arg: Base64 string to be decoded. Unicode input will be
-                encoded to UTF-8 before Base64 decoding.
-    :type arg: str or unicode
-
-    :returns: Decoded string.
-    :rtype: str
+    :returns: Decoded data. Unicode if `decoding` is not None.
+    :rtype: str or unicode
 
     """
-    normalized = _to_utf8(arg)
-    return base64.urlsafe_b64decode(
-        normalized + '=' * (4 - (len(normalized) % 4)))
+    decoded = base64.urlsafe_b64decode(arg + '=' * (4 - (len(arg) % 4)))
+    return decoded if decoding is None else decoded.decode(decoding)
diff --git a/letsencrypt/client/le_util_test.py b/letsencrypt/client/le_util_test.py
index 66c21b473..b5126218b 100644
--- a/letsencrypt/client/le_util_test.py
+++ b/letsencrypt/client/le_util_test.py
@@ -72,32 +72,32 @@ class CheckPermissionsTest(unittest.TestCase):
 class JOSEB64EncodeTest(unittest.TestCase):
     """Tests for letsencrypt.client.le_util.jose_b64encode."""
 
-    def _call(self, arg):
+    def _call(self, data, encoding):
         from letsencrypt.client.le_util import jose_b64encode
-        return jose_b64encode(arg)
+        return jose_b64encode(data, encoding)
 
-    def test_str(self):
-        self.assertEqual(self._call('foo'), 'Zm9v')
+    def test_without_encoding(self):
+        self.assertEqual(self._call('foo', None), 'Zm9v')
 
-    def test_unicode(self):
-        self.assertEqual(self._call(u'\u0105'), 'xIU')
+    def test_with_encoding(self):
+        self.assertEqual(self._call(u'\u0105', 'utf-8'), 'xIU')
 
 
 class JOSEB64DecodeTest(unittest.TestCase):
     """Tests for letsencrypt.client.le_util.jose_b64decode."""
 
-    def _call(self, arg):
+    def _call(self, jose_b64_string, decoding):
         from letsencrypt.client.le_util import jose_b64decode
-        return jose_b64decode(arg)
+        return jose_b64decode(jose_b64_string, decoding)
 
-    def test_str(self):
-        self.assertEqual(self._call('Zm9v='), 'foo')
+    def test_without_decoding(self):
+        self.assertEqual(self._call('Zm9v=', None), 'foo')
 
-    def test_unicode(self):
-        self.assertEqual(self._call(u'XIU='), '\\\x85')
+    def test_with_encoding(self):
+        self.assertEqual(self._call('xIU=', 'utf-8'), u'\u0105')
 
     def test_fills_padding(self):
-        self.assertEqual(self._call('Zm9v'), 'foo')
+        self.assertEqual(self._call('Zm9v', None), 'foo')
 
 
 if __name__ == '__main__':