Merge pull request #1 from matrix-org/daniel/urlsafe

Use url-safe base64 charset

test_unpaddedbase64.py

@@ -31,3 +31,11 @@ def test_decode(self):
         self.assertEqual(decode_base64(u'AAAA'), b'\x00\x00\x00')
         with self.assertRaises(Exception):
             decode_base64(u'A')
+
+    def test_encode_urlunsafe_chars(self):
+        self.assertEqual(encode_base64(b'\xff\xe6\x9a'), u'/+aa')
+        self.assertEqual(encode_base64(b'\xff\xe6\x9a', True), u'_-aa')
+
+    def test_decode_urlunsafe_chars(self):
+        self.assertEqual(decode_base64(u'/+aa'), b'\xff\xe6\x9a')
+        self.assertEqual(decode_base64(u'_-aa'), b'\xff\xe6\x9a')

unpaddedbase64.py

@@ -17,10 +17,11 @@
 __version__ = "1.0.1"
 
 
-def encode_base64(input_bytes):
+def encode_base64(input_bytes, urlsafe=False):
     """Encode bytes as a base64 string without any padding."""
 
-    output_bytes = base64.b64encode(input_bytes)
+    encode = base64.urlsafe_b64encode if urlsafe else base64.b64encode
+    output_bytes = encode(input_bytes)
     output_string = output_bytes.decode("ascii")
     return output_string.rstrip(u"=")
 
@@ -32,5 +33,8 @@ def decode_base64(input_string):
     input_bytes = input_string.encode("ascii")
     input_len = len(input_bytes)
     padding = b"=" * (3 - ((input_len + 3) % 4))
-    output_bytes = base64.b64decode(input_bytes + padding)
+    decode = base64.b64decode
+    if u'-' in input_string or u'_' in input_string:
+        decode = base64.urlsafe_b64decode
+    output_bytes = decode(input_bytes + padding)
     return output_bytes