Update auth_tkt.py for python3 compatibility (#45)

* Update auth_tkt.py urllib imports to work in python3

* Add tests for auth AuthTicket

* Adapt auth_tkt to be python2 and python3 compatible

2 files changed, 150 insertions, 20 deletions

diff --git a/paste/auth/auth_tkt.py b/paste/auth/auth_tkt.py
index da8ddbd..24fc2cb 100644
--- a/paste/auth/auth_tkt.py
+++ b/paste/auth/auth_tkt.py
@@ -36,7 +36,7 @@ it's primary benefit is compatibility with mod_auth_tkt, which in turn
 makes it possible to use the same authentication process with
 non-Python code run under Apache.
 """
-
+import six
 import time as time_mod
 try:
     import hashlib
@@ -49,8 +49,13 @@ except ImportError:
     # Python 2
     from Cookie import SimpleCookie
 from paste import request
-from urllib import quote as url_quote
-from urllib import unquote as url_unquote
+
+try:
+    from urllib import quote as url_quote # Python 2.X
+    from urllib import unquote as url_unquote
+except ImportError:
+    from urllib.parse import quote as url_quote  # Python 3+
+    from urllib.parse import unquote as url_unquote
 
 DEFAULT_DIGEST = hashlib.md5
 
@@ -98,7 +103,7 @@ class AuthTicket(object):
         self.secret = secret
         self.userid = userid
         self.ip = ip
-        if not isinstance(tokens, basestring):
+        if not isinstance(tokens, six.string_types):
             tokens = ','.join(tokens)
         self.tokens = tokens
         self.user_data = user_data
@@ -108,7 +113,7 @@ class AuthTicket(object):
             self.time = time
         self.cookie_name = cookie_name
         self.secure = secure
-        if isinstance(digest_algo, str):
+        if isinstance(digest_algo, six.binary_type):
             # correct specification of digest from hashlib or fail
             self.digest_algo = getattr(hashlib, digest_algo)
         else:
@@ -120,15 +125,20 @@ class AuthTicket(object):
             self.user_data, self.digest_algo)
 
     def cookie_value(self):
-        v = '%s%08x%s!' % (self.digest(), int(self.time), url_quote(self.userid))
+        v = b'%s%08x%s!' % (self.digest(), int(self.time), maybe_encode(url_quote(self.userid)))
         if self.tokens:
-            v += self.tokens + '!'
-        v += self.user_data
+            v += maybe_encode(self.tokens) + b'!'
+        v += maybe_encode(self.user_data)
         return v
 
     def cookie(self):
         c = SimpleCookie()
-        c[self.cookie_name] = self.cookie_value().encode('base64').strip().replace('\n', '')
+        if six.PY3:
+            import base64
+            cookie_value = base64.b64encode(self.cookie_value())
+        else:
+            cookie_value = self.cookie_value().encode('base64').strip().replace('\n', '')
+        c[self.cookie_name] = cookie_value
         c[self.cookie_name]['path'] = '/'
         if self.secure:
             c[self.cookie_name]['secure'] = 'true'
@@ -154,7 +164,7 @@ def parse_ticket(secret, ticket, ip, digest_algo=DEFAULT_DIGEST):
     If the ticket cannot be parsed, ``BadTicket`` will be raised with
     an explanation.
     """
-    if isinstance(digest_algo, str):
+    if isinstance(digest_algo, six.binary_type):
         # correct specification of digest from hashlib or fail
         digest_algo = getattr(hashlib, digest_algo)
     digest_hexa_size = digest_algo().digest_size * 2
@@ -196,26 +206,26 @@ def calculate_digest(ip, timestamp, secret, userid, tokens, user_data,
     userid = maybe_encode(userid)
     tokens = maybe_encode(tokens)
     user_data = maybe_encode(user_data)
-    digest0 = digest_algo(
-        encode_ip_timestamp(ip, timestamp) + secret + userid + '\0'
-        + tokens + '\0' + user_data).hexdigest()
+    digest0 = maybe_encode(digest_algo(
+        encode_ip_timestamp(ip, timestamp) + secret + userid + b'\0'
+        + tokens + b'\0' + user_data).hexdigest())
     digest = digest_algo(digest0 + secret).hexdigest()
-    return digest
+    return maybe_encode(digest)
 
 
 def encode_ip_timestamp(ip, timestamp):
-    ip_chars = ''.join(map(chr, map(int, ip.split('.'))))
+    ip_chars = b''.join(map(six.int2byte, map(int, ip.split('.'))))
     t = int(timestamp)
     ts = ((t & 0xff000000) >> 24,
           (t & 0xff0000) >> 16,
           (t & 0xff00) >> 8,
           t & 0xff)
-    ts_chars = ''.join(map(chr, ts))
-    return ip_chars + ts_chars
+    ts_chars = b''.join(map(six.int2byte, ts))
+    return (ip_chars + ts_chars)
 
 
 def maybe_encode(s, encoding='utf8'):
-    if isinstance(s, unicode):
+    if isinstance(s, six.text_type):
         s = s.encode(encoding)
     return s
 
@@ -349,7 +359,7 @@ class AuthTKTMiddleware(object):
         return self.app(environ, cookie_setting_start_response)
 
     def set_user_cookie(self, environ, userid, tokens, user_data):
-        if not isinstance(tokens, basestring):
+        if not isinstance(tokens, six.string_types):
             tokens = ','.join(tokens)
         if self.include_ip:
             remote_addr = environ['REMOTE_ADDR']
@@ -415,7 +425,7 @@ def make_auth_tkt_middleware(
     Creates the `AuthTKTMiddleware
     <class-paste.auth.auth_tkt.AuthTKTMiddleware.html>`_.
 
-    ``secret`` is requird, but can be set globally or locally.
+    ``secret`` is required, but can be set globally or locally.
     """
     from paste.deploy.converters import asbool
     secure = asbool(secure)
diff --git a/tests/test_auth/test_auth_tkt.py b/tests/test_auth/test_auth_tkt.py
new file mode 100644
index 0000000..4713ea0
--- /dev/null
+++ b/tests/test_auth/test_auth_tkt.py
@@ -0,0 +1,120 @@
+import hashlib
+import six
+if six.PY3:
+    import base64
+from paste.auth.auth_tkt import AuthTicket
+try:
+    from http.cookies import SimpleCookie
+except ImportError:
+    # Python 2
+    from Cookie import SimpleCookie
+
+
+def test_auth_ticket_digest_and_cookie_value():
+    test_parameters = [
+        (
+            (
+                'shared_secret',
+                'username',
+                '0.0.0.0',  # remote address
+            ),
+            {
+                'tokens': ['admin'],
+                'time': 1579782607
+            },
+            b'731274bec45f6983c1f33bac8e8baf43',
+            b'731274bec45f6983c1f33bac8e8baf435e2991cfusername!admin!',
+        ),
+        (
+            (
+                'shared_secret',
+                'username',
+                '0.0.0.0',
+            ),
+            {
+                'tokens': ['admin'],
+                'time': 1579782607,
+                'digest_algo': hashlib.sha512
+            },
+            b'09e72a63c57ca4cfeca5fa578646deb2b27f7a461d91ad9aa32b85c93ef6fa7744ac006eb3d9a71a36375b5ab50cbae072bb3042e2a59198b7f314900cba4423',
+            b'09e72a63c57ca4cfeca5fa578646deb2b27f7a461d91ad9aa32b85c93ef6fa7744ac006eb3d9a71a36375b5ab50cbae072bb3042e2a59198b7f314900cba44235e2991cfusername!admin!',
+        ),
+    ]
+
+    for test_args, test_kwargs, expected_digest, expected_cookie_value in test_parameters:
+        token = AuthTicket(*test_args, **test_kwargs)
+        assert expected_digest == token.digest()
+        assert expected_cookie_value == token.cookie_value()
+
+
+def test_auth_ticket_cookie():
+    test_parameters = [
+        (
+            (
+                'shared_secret',
+                'username',
+                '0.0.0.0',  # remote address
+            ),
+            {
+                'tokens': ['admin'],
+                'time': 1579782607
+            },
+            {
+                'name': 'auth_tkt',
+                'path': '/',
+                'secure': '',
+                'cookie_value': b'731274bec45f6983c1f33bac8e8baf435e2991cfusername!admin!'
+            }
+        ),
+        (
+            (
+                'shared_secret',
+                'username',
+                '0.0.0.0',  # remote address
+            ),
+            {
+                'tokens': ['admin'],
+                'time': 1579782607,
+                'secure': True
+            },
+            {
+                'name': 'auth_tkt',
+                'path': '/',
+                'secure': 'true',
+                'cookie_value': b'731274bec45f6983c1f33bac8e8baf435e2991cfusername!admin!'
+            }
+        ),
+        (
+            (
+                'shared_secret',
+                'username',
+                '0.0.0.0',  # remote address
+            ),
+            {
+                'tokens': ['admin'],
+                'time': 1579782607,
+                'cookie_name': 'custom_cookie',
+                'secure': False
+            },
+            {
+                'name': 'custom_cookie',
+                'path': '/',
+                'secure': '',
+                'cookie_value': b'731274bec45f6983c1f33bac8e8baf435e2991cfusername!admin!'
+            }
+        ),
+    ]
+
+    for test_args, test_kwargs, expected_values in test_parameters:
+        token = AuthTicket(*test_args, **test_kwargs)
+        expected_cookie = SimpleCookie()
+        if six.PY3:
+            # import pdb; pdb.set_trace()
+            expected_cookie_value = base64.b64encode(expected_values['cookie_value'])
+        else:
+            expected_cookie_value = expected_values['cookie_value'].encode('base64')
+
+        expected_cookie[expected_values['name']] = expected_cookie_value
+        expected_cookie[expected_values['name']]['path'] = expected_values['path']
+        expected_cookie[expected_values['name']]['secure'] = expected_values['secure']
+        assert expected_cookie == token.cookie()