Merge pull request #42 from rebeckag/unicode_cleanup

Unicode with base64

Author: Roland Hedberg

src/jwkest/__init__.py

@@ -3,6 +3,7 @@
 import logging
 import re
 import struct
+import six
 
 try:
     from builtins import zip
@@ -103,12 +104,13 @@ def long_to_base64(n):
     if not len(data):
         data = '\x00'
     s = base64.urlsafe_b64encode(data).rstrip(b'=')
-    return s
+    return s.decode("ascii")
 
 
 def base64_to_long(data):
-    # if isinstance(data, str):
-    #     data = bytes(data)
+    if isinstance(data, six.text_type):
+        data = data.encode("ascii")
+
     # urlsafe_b64decode will happily convert b64encoded data
     _d = base64.urlsafe_b64decode(bytes(data) + b'==')
     return intarr2long(struct.unpack('%sB' % len(_d), _d))
@@ -164,10 +166,7 @@ def b64d(b):
     :param b: bytes
     """
 
-    if b.endswith(b'='):  # shouldn't but there you are
-        cb = b.split(b'=')[0]
-    else:
-        cb = b
+    cb = b.rstrip(b"=")  # shouldn't but there you are
 
     # Python's base64 functions ignore invalid characters, so we need to
     # check for them explicitly.

tests/test_0_jwkest.py

@@ -1,6 +1,7 @@
+import base64
 import os
 import struct
-from jwkest import long2intarr
+from jwkest import long2intarr, b64d, b64e
 from jwkest import intarr2long
 from jwkest import base64_to_long
 from jwkest import long_to_base64
@@ -14,11 +15,13 @@
 def full_path(local_file):
     return os.path.join(BASEDIR, local_file)
 
+
 CERT = full_path("cert.pem")
 KEY = full_path("server.key")
 
 _CKEY = pem_cert2rsa(CERT)
 
+
 def test_long_intarr_long():
     ia = long2intarr(_CKEY.n)
     _n = intarr2long(ia)
@@ -39,5 +42,16 @@ def test_long_base64_long():
     assert _CKEY.n == l
 
 
+def test_b64d_with_padded_data():
+    data = "abcd".encode("utf-8")
+    encoded = base64.urlsafe_b64encode(data)
+    assert b64d(encoded) == data
+
+
+def test_b64_encode_decode():
+    data = "abcd".encode("utf-8")
+    assert b64d(b64e(data)) == data
+
+
 if __name__ == "__main__":
-    test_long_base64_long()
+    test_long_base64_long()

tests/test_2_jwk.py

@@ -19,6 +19,7 @@
 from jwkest.jwk import RSAKey
 from jwkest.jwk import base64_to_long
 import os.path
+import pytest
 
 __author__ = 'rohe0002'
 BASEDIR = os.path.abspath(os.path.dirname(__file__))
@@ -31,12 +32,12 @@ def full_path(local_file):
 CERT = full_path("cert.pem")
 KEY = full_path("server.key")
 
-N = b'wf-wiusGhA-gleZYQAOPQlNUIucPiqXdPVyieDqQbXXOPBe3nuggtVzeq7pVFH1dZz4dY2Q2LA5DaegvP8kRvoSB_87ds3dy3Rfym_GUSc5B0l1TgEobcyaep8jguRoHto6GWHfCfKqoUYZq4N8vh4LLMQwLR6zi6Jtu82nB5k8'
-E = b'AQAB'
+N = 'wf-wiusGhA-gleZYQAOPQlNUIucPiqXdPVyieDqQbXXOPBe3nuggtVzeq7pVFH1dZz4dY2Q2LA5DaegvP8kRvoSB_87ds3dy3Rfym_GUSc5B0l1TgEobcyaep8jguRoHto6GWHfCfKqoUYZq4N8vh4LLMQwLR6zi6Jtu82nB5k8'
+E = 'AQAB'
 
 JWK = {"keys": [
-    {'kty': 'RSA', 'use': 'foo', 'e': E.decode("utf-8"), 'kid': "abc",
-     'n': N.decode("utf8")}
+    {'kty': 'RSA', 'use': 'foo', 'e': E, 'kid': "abc",
+     'n': N}
 ]}
 
 
@@ -54,14 +55,12 @@ def test_urlsafe_base64decode():
     s0 = base64.b64encode(data)
     # try to convert it back to long, should throw an exception if the strict
     # function is used
-    try:
-        l = base64url_to_long(s0)
-    except ValueError:
-        pass
-    else:
-        assert False
-    # Not else
+    with pytest.raises(ValueError):
+        base64url_to_long(s0)
+
+    # Not else, should not raise exception
     l = base64_to_long(s0)
+    assert l
 
 
 def test_pem_cert2rsa():
@@ -87,8 +86,8 @@ def test_kspec():
     print(_key)
     jwk = _key.serialize()
     assert jwk["kty"] == "RSA"
-    assert jwk["e"] == JWK["keys"][0]["e"].encode("utf-8")
-    assert jwk["n"] == JWK["keys"][0]["n"].encode("utf-8")
+    assert jwk["e"] == JWK["keys"][0]["e"]
+    assert jwk["n"] == JWK["keys"][0]["n"]
 
 
 def test_loads_0():
@@ -170,9 +169,8 @@ def test_import_rsa_key():
     djwk = jwk_wrap(_ckey).to_dict()
     print(djwk)
     assert _eq(djwk.keys(), ["kty", "e", "n", "p", "q", "d"])
-    assert djwk[
-               "n"] == b'5zbNbHIYIkGGJ3RGdRKkYmF4gOorv5eDuUKTVtuu3VvxrpOWvwnFV-NY0LgqkQSMMyVzodJE3SUuwQTUHPXXY5784vnkFqzPRx6bHgPxKz7XfwQjEBTafQTMmOeYI8wFIOIHY5i0RWR-gxDbh_D5TXuUqScOOqR47vSpIbUH-nc'
-    assert djwk['e'] == b'AQAB'
+    assert djwk["n"] == '5zbNbHIYIkGGJ3RGdRKkYmF4gOorv5eDuUKTVtuu3VvxrpOWvwnFV-NY0LgqkQSMMyVzodJE3SUuwQTUHPXXY5784vnkFqzPRx6bHgPxKz7XfwQjEBTafQTMmOeYI8wFIOIHY5i0RWR-gxDbh_D5TXuUqScOOqR47vSpIbUH-nc'
+    assert djwk['e'] == 'AQAB'
 
 
 def test_serialize_rsa_pub_key():
@@ -238,25 +236,15 @@ def test_cmp_rsa_ec():
 
     _key2 = ECKey(**ECKEY)
 
-    try:
-        assert _key1 == _key2
-    except AssertionError:
-        pass
-    else:
-        assert False
+    assert _key1 != _key2
 
 
 def test_cmp_neq_ec():
     priv, pub = P256.key_pair()
     _key1 = ECKey(x=pub[0], y=pub[1], d=priv, crv="P-256")
     _key2 = ECKey(**ECKEY)
 
-    try:
-        assert _key1 == _key2
-    except AssertionError:
-        pass
-    else:
-        assert False
+    assert _key1 != _key2
 
 
 JWKS = {"keys": [

tests/test_3_jws.py

@@ -26,6 +26,7 @@
 def full_path(local_file):
     return os.path.join(BASEDIR, local_file)
 
+
 KEY = full_path("server.key")
 
 JWK = {"keys": [{'alg': 'RSA',
@@ -44,7 +45,6 @@ def full_path(local_file):
             119, 98, 61, 34, 61, 46, 33, 114, 5, 46, 79, 8, 192, 205, 154, 245,
             103, 208, 128, 163]
 
-
 JWKS = {"keys": [
     {
         "n": b"zkpUgEgXICI54blf6iWiD2RbMDCOO1jV0VSff1MFFnujM4othfMsad7H1kRo50YM5S_X9TdvrpdOfpz5aBaKFhT6Ziv0nhtcekq1eRl8mjBlvGKCE5XGk-0LFSDwvqgkJoFYInq7bu0a4JEzKs5AyJY75YlGh879k1Uu2Sv3ZZOunfV1O1Orta-NvS-aG_jN5cstVbCGWE20H0vFVrJKNx0Zf-u-aA-syM4uX7wdWgQ-owoEMHge0GmGgzso2lwOYf_4znanLwEuO3p5aabEaFoKNR4K6GjQcjBcYmDEE4CtfRU9AEmhcD1kleiTB9TjPWkgDmT9MXsGxBHf3AKT5w",
@@ -132,6 +132,7 @@ def full_path(local_file):
 SIGKEYS = KEYS()
 SIGKEYS.load_dict(JWKS)
 
+
 def test_1():
     claimset = {"iss": "joe",
                 "exp": 1300819380,
@@ -148,14 +149,14 @@ def test_1():
 
 
 def test_hmac_256():
-    payload = b'Please take a moment to register today'
+    payload = 'Please take a moment to register today'
     keys = [SYMKey(key=jwkest.intarr2bin(HMAC_KEY))]
     _jws = JWS(payload, alg="HS256")
     _jwt = _jws.sign_compact(keys)
 
     info = JWS().verify_compact(_jwt, keys)
 
-    assert info == payload.decode("utf-8")
+    assert info == payload
 
 
 def test_hmac_384():
@@ -205,7 +206,7 @@ def test_left_hash_hs512():
 def test_rs256():
     payload = "Please take a moment to register today"
     keys = [RSAKey(key=import_rsa_key_from_file(KEY))]
-    #keys[0]._keytype = "private"
+    # keys[0]._keytype = "private"
     _jws = JWS(payload, alg="RS256")
     _jwt = _jws.sign_compact(keys)
 
@@ -218,7 +219,7 @@ def test_rs256():
 def test_rs384():
     payload = "Please take a moment to register today"
     keys = [RSAKey(key=import_rsa_key_from_file(KEY))]
-    #keys[0]._keytype = "private"
+    # keys[0]._keytype = "private"
     _jws = JWS(payload, alg="RS384")
     _jwt = _jws.sign_compact(keys)
 
@@ -230,7 +231,7 @@ def test_rs384():
 def test_rs512():
     payload = "Please take a moment to register today"
     keys = [RSAKey(key=import_rsa_key_from_file(KEY))]
-    #keys[0]._keytype = "private"
+    # keys[0]._keytype = "private"
     _jws = JWS(payload, alg="RS512")
     _jwt = _jws.sign_compact(keys)
 
@@ -323,7 +324,7 @@ def test_signer_es512():
     payload = "Please take a moment to register today"
     _key = ECKey().load_key(P521)
     keys = [_key]
-    #keys[0]._keytype = "private"
+    # keys[0]._keytype = "private"
     _jws = JWS(payload, alg="ES512")
     _jwt = _jws.sign_compact(keys)
 
@@ -335,7 +336,7 @@ def test_signer_es512():
 def test_signer_ps256():
     payload = "Please take a moment to register today"
     keys = [RSAKey(key=import_rsa_key_from_file(KEY))]
-    #keys[0]._keytype = "private"
+    # keys[0]._keytype = "private"
     _jws = JWS(payload, alg="PS256")
     _jwt = _jws.sign_compact(keys)
 
@@ -347,7 +348,7 @@ def test_signer_ps256():
 def test_signer_ps256_fail():
     payload = "Please take a moment to register today"
     keys = [RSAKey(key=import_rsa_key_from_file(KEY))]
-    #keys[0]._keytype = "private"
+    # keys[0]._keytype = "private"
     _jws = JWS(payload, alg="PS256")
     _jwt = _jws.sign_compact(keys)[:-5] + 'abcde'
 
@@ -363,7 +364,7 @@ def test_signer_ps256_fail():
 def test_signer_ps384():
     payload = "Please take a moment to register today"
     keys = [RSAKey(key=import_rsa_key_from_file(KEY))]
-    #keys[0]._keytype = "private"
+    # keys[0]._keytype = "private"
     _jws = JWS(payload, alg="PS384")
     _jwt = _jws.sign_compact(keys)
 
@@ -376,7 +377,7 @@ def test_signer_ps512():
     payload = "Please take a moment to register today"
     # Key has to be big enough  > 512+512+2
     keys = [RSAKey(key=import_rsa_key_from_file(full_path("./size2048.key")))]
-    #keys[0]._keytype = "private"
+    # keys[0]._keytype = "private"
     _jws = JWS(payload, alg="PS512")
     _jwt = _jws.sign_compact(keys)
 
@@ -410,20 +411,21 @@ def test_sign_2():
          "use": "sig",
          "kid": "af22448d-4c7b-464d-b63a-f5bd90f6d7d1",
          "n": "o9g8DpUwBW6B1qmcm-TfEh4rNX7n1t38jdo4Gkl_cI3q--7n0Blg0kN88LHZvyZjUB2NhBdFYNxMP8ucy0dOXvWGWzaPmGnq3DM__lN8P4WjD1cCTAVEYKawNBAmGKqrFj1SgpPNsSqiqK-ALM1w6mZ-QGimjOgwCyJy3l9lzZh5D8tKnS2t1pZgE0X5P7lZQWHYpHPqp4jKhETzrCpPGfv0Rl6nmmjp7NlRYBkWKf_HEKE333J6M039m2FbKgxrBg3zmYYpmHuMzVgxxb8LSiv5aqyeyJjxM-YDUAgNQBfKNhONqXyu9DqtSprNkw6sqmuxK0QUVrNYl3b03PgS5Q"
-        }]}
+         }]}
 
     keys = KEYS()
     keys.load_dict(keyset)
     jws = JWS("payload", alg="RS512")
     jws.sign_compact(keys=keys)
 
+
 def test_signer_protected_headers():
     payload = "Please take a moment to register today"
     _key = ECKey().load_key(P256)
     keys = [_key]
     _jws = JWS(payload, alg="ES256")
     protected = dict(header1=u"header1 is protected",
-        header2="header2 is protected too", a=1)
+                     header2="header2 is protected too", a=1)
     _jwt = _jws.sign_compact(keys, protected=protected)
 
     exp_protected = protected.copy()
@@ -436,32 +438,26 @@ def test_signer_protected_headers():
     info = _rj.verify_compact(_jwt, keys)
     assert info == payload
 
+
 def test_verify_protected_headers():
     payload = "Please take a moment to register today"
     _key = ECKey().load_key(P256)
     keys = [_key]
     _jws = JWS(payload, alg="ES256")
     protected = dict(header1=u"header1 is protected",
-        header2="header2 is protected too", a=1)
+                     header2="header2 is protected too", a=1)
     _jwt = _jws.sign_compact(keys, protected=protected)
     protectedHeader, enc_payload, sig = _jwt.split(".")
     data = dict(payload=enc_payload, signatures=[
         dict(
             header=dict(alg=u"ES256", jwk=_key.serialize()),
             protected=protectedHeader,
             signature=sig,
-            )
-        ])
-    fobj = io.BytesIO(JSONEncoder().encode(data).encode("utf-8"))
+        )
+    ])
+    stream = io.StringIO(json.dumps(data, ensure_ascii=False))
     _jws = JWS()
-    reader = codecs.getreader("utf-8")
-    assert _jws.verify_json(reader(fobj)) == payload
-
-class JSONEncoder(json.JSONEncoder):
-    def default(self, o):
-        if isinstance(o, bytes):
-            return o.decode('utf-8')
-        return json.JSONEncoder.default(self, o)
+    assert _jws.verify_json(stream) == payload
 
 
 def test_pick():