From aef18c4413c4d480eef786a9ae47efd024e57da7 Mon Sep 17 00:00:00 2001
From: Jakub Warmuz <jakub@warmuz.org>
Date: Mon, 24 Nov 2014 11:00:08 +0100
Subject: [PATCH] JOSE Base64: encode str only, decode str or ascii unicode

---
 letsencrypt/client/le_util.py      | 41 ++++++++++++++++--------------
 letsencrypt/client/le_util_test.py | 30 +++++++++++-----------
 2 files changed, 37 insertions(+), 34 deletions(-)

diff --git a/letsencrypt/client/le_util.py b/letsencrypt/client/le_util.py
index 8a87d7c09..19070858f 100644
--- a/letsencrypt/client/le_util.py
+++ b/letsencrypt/client/le_util.py
@@ -74,39 +74,42 @@ def unique_file(default_name, mode=0777):
 #   - padding stripped
 
 
-def jose_b64encode(data, encoding='utf-8'):
+def jose_b64encode(data):
     """JOSE Base64 encode.
 
     :param data: Data to be encoded.
-    :type data: str or unicode
+    :type data: str or bytearray
 
-    :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
+    :raises: TypeError
 
     :returns: JOSE Base64 string.
     :rtype: str
 
     """
-    encoded = data if encoding is None else data.encode(encoding)
-    return base64.urlsafe_b64encode(encoded).rstrip('=')
+    if not isinstance(data, str):
+        raise TypeError('argument should be str or bytearray')
+    return base64.urlsafe_b64encode(data).rstrip('=')
 
 
-def jose_b64decode(arg, decoding='utf-8'):
+def jose_b64decode(data):
     """JOSE Base64 decode.
 
-    :param arg: Base64 string to be decoded.
-    :type arg: str or unicode
+    :param data: Base64 string to be decoded. If it's unicode, then
+                 only ASCII characters are allowed.
+    :type data: str or unicode
 
-    :param decoding: Name of the encoding to be performed after
-                     Base64 decoding.
-    :type decoding: str or None
+    :raises: ValueError, TypeError
 
-    :returns: Decoded data. Unicode if `decoding` is not None.
-    :rtype: str or unicode
+    :returns: Decoded data.
 
     """
-    ascii = arg.encode('ascii')  # equivalent to str(arg), for unicode input
-    decoded = base64.urlsafe_b64decode(ascii + '=' * (4 - (len(ascii) % 4)))
-    return decoded if decoding is None else decoded.decode(decoding)
+    if isinstance(data, unicode):
+        try:
+            data = data.encode('ascii')
+        except UnicodeEncodeError:
+            raise ValueError(
+                'unicode argument should contain only ASCII characters')
+    elif not isinstance(data, str):
+        raise TypeError('argument should be a str or unicode')
+
+    return base64.urlsafe_b64decode(data + '=' * (4 - (len(data) % 4)))
diff --git a/letsencrypt/client/le_util_test.py b/letsencrypt/client/le_util_test.py
index 6c2714a8d..30743c24a 100644
--- a/letsencrypt/client/le_util_test.py
+++ b/letsencrypt/client/le_util_test.py
@@ -88,46 +88,46 @@ B64_URL_UNSAFE_EXAMPLES = {
 class JOSEB64EncodeTest(unittest.TestCase):
     """Tests for letsencrypt.client.le_util.jose_b64encode."""
 
-    def _call(self, data, encoding):
+    def _call(self, data):
         from letsencrypt.client.le_util import jose_b64encode
-        return jose_b64encode(data, encoding)
+        return jose_b64encode(data)
 
     def test_unsafe_url(self):
         for text, b64 in B64_URL_UNSAFE_EXAMPLES.iteritems():
-            self.assertEqual(self._call(text, None), b64)
+            self.assertEqual(self._call(text), b64)
 
     def test_different_paddings(self):
         for text, (b64, _) in JOSE_B64_PADDING_EXAMPLES.iteritems():
-            self.assertEqual(self._call(text, None), b64)
+            self.assertEqual(self._call(text), b64)
 
-    def test_with_encoding(self):
-        self.assertEqual(self._call(u'\u0105', 'utf-8'), 'xIU')
+    def test_unicode_fails_with_type_error(self):
+        self.assertRaises(TypeError, self._call, u'some unicode')
 
 
 class JOSEB64DecodeTest(unittest.TestCase):
     """Tests for letsencrypt.client.le_util.jose_b64decode."""
 
-    def _call(self, jose_b64_string, decoding):
+    def _call(self, data):
         from letsencrypt.client.le_util import jose_b64decode
-        return jose_b64decode(jose_b64_string, decoding)
+        return jose_b64decode(data)
 
     def test_unsafe_url(self):
         for text, b64 in B64_URL_UNSAFE_EXAMPLES.iteritems():
-            self.assertEqual(self._call(b64, None), text)
+            self.assertEqual(self._call(b64), text)
 
     def test_input_without_padding(self):
         for text, (b64, _) in JOSE_B64_PADDING_EXAMPLES.iteritems():
-            self.assertEqual(self._call(b64, None), text)
+            self.assertEqual(self._call(b64), text)
 
     def test_input_with_padding(self):
         for text, (b64, pad) in JOSE_B64_PADDING_EXAMPLES.iteritems():
-            self.assertEqual(self._call(b64 + pad, None), text)
+            self.assertEqual(self._call(b64 + pad), text)
 
-    def test_with_encoding(self):
-        self.assertEqual(self._call('xIU=', 'utf-8'), u'\u0105')
+    def test_unicode_with_ascii(self):
+        self.assertEqual(self._call(u'YQ'), 'a')
 
-    def test_unicode(self):
-        self.assertEqual(self._call(u'YQ', None), 'a')
+    def test_non_ascii_unicode_fails(self):
+        self.assertRaises(ValueError, self._call, u'\u0105')
 
 
 if __name__ == '__main__':