From 7483ece150564b242ba0ac5c091319ee570dd9e1 Mon Sep 17 00:00:00 2001
From: David Tulloh <git-david@tulloh.id.au>
Date: Tue, 31 May 2016 18:24:04 +0200
Subject: [PATCH] Added alternate hash support to auth_tkt plugin

---
 repoze/who/plugins/tests/test_authtkt.py |  124 ++++++++++++++++++++++++++++++++++++++++-
 repoze/who/plugins/auth_tkt.py           |   25 +++++++-
 2 files changed, 142 insertions(+), 7 deletions(-)

diff --git a/repoze/who/plugins/auth_tkt.py b/repoze/who/plugins/auth_tkt.py
index 1c551b1..db74da5 100644
--- a/repoze/who/plugins/auth_tkt.py
+++ b/repoze/who/plugins/auth_tkt.py
@@ -5,6 +5,10 @@
 from codecs import utf_8_encode
 import os
 import time
+try:
+    import hashlib
+except ImportError:
+    import md5 as hashlib # Will only support md5 algorithm
 from wsgiref.handlers import _monthname     # Locale-independent, RFC-2616
 from wsgiref.handlers import _weekdayname   # Locale-independent, RFC-2616
 try:
@@ -28,6 +32,7 @@
         return _UTCNOW
     return datetime.datetime.utcnow()
 
+DEFAULT_DIGEST = hashlib.md5
 
 @implementer(IIdentifier, IAuthenticator)
 class AuthTktCookiePlugin(object):
@@ -51,7 +56,8 @@
  
     def __init__(self, secret, cookie_name='auth_tkt',
                  secure=False, include_ip=False,
-                 timeout=None, reissue_time=None, userid_checker=None):
+                 timeout=None, reissue_time=None, userid_checker=None,
+                 digest_algo=DEFAULT_DIGEST):
         self.secret = secret
         self.cookie_name = cookie_name
         self.include_ip = include_ip
@@ -62,6 +68,7 @@
         self.timeout = timeout
         self.reissue_time = reissue_time
         self.userid_checker = userid_checker
+        self.digest_algo = digest_algo
 
     # IIdentifier
     def identify(self, environ):
@@ -78,7 +85,7 @@
         
         try:
             timestamp, userid, tokens, user_data = auth_tkt.parse_ticket(
-                self.secret, cookie.value, remote_addr)
+                self.secret, cookie.value, remote_addr, self.digest_algo)
         except auth_tkt.BadTicket:
             return None
 
@@ -126,7 +133,8 @@
         if old_cookie_value:
             try:
                 timestamp,userid,tokens,userdata = auth_tkt.parse_ticket(
-                    self.secret, old_cookie_value, remote_addr)
+                    self.secret, old_cookie_value, remote_addr,
+                    self.digest_algo)
             except auth_tkt.BadTicket:
                 pass
         tokens = tuple(tokens)
@@ -155,7 +163,8 @@
                 tokens=who_tokens,
                 user_data=who_userdata,
                 cookie_name=self.cookie_name,
-                secure=self.secure)
+                secure=self.secure,
+                digest_algo=self.digest_algo)
             new_cookie_value = ticket.cookie_value()
             
             if old_cookie_value != new_cookie_value:
@@ -226,6 +235,7 @@
                 timeout=None,
                 reissue_time=None,
                 userid_checker=None,
+                digest_algo=DEFAULT_DIGEST,
                ):
     from repoze.who.utils import resolveDotted
     if (secret is None and secretfile is None):
@@ -244,6 +254,12 @@
         reissue_time = int(reissue_time)
     if userid_checker is not None:
         userid_checker = resolveDotted(userid_checker)
+    if isinstance(digest_algo, str):
+        try:
+            digest_algo = getattr(hashlib, digest_algo)
+        except AttributeError:
+            raise ValueError("No such 'digest_algo': %s" % digest_algo)
+
     plugin = AuthTktCookiePlugin(secret,
                                  cookie_name,
                                  _bool(secure),
@@ -251,6 +267,7 @@
                                  timeout,
                                  reissue_time,
                                  userid_checker,
+                                 digest_algo,
                                  )
     return plugin
 
diff --git a/repoze/who/plugins/tests/test_authtkt.py b/repoze/who/plugins/tests/test_authtkt.py
index b01c6c9..accf88a 100644
--- a/repoze/who/plugins/tests/test_authtkt.py
+++ b/repoze/who/plugins/tests/test_authtkt.py
@@ -36,8 +36,7 @@
     def _makeTicket(self, userid='userid', remote_addr='0.0.0.0',
                     tokens = [], userdata='userdata',
                     cookie_name='auth_tkt', secure=False,
-                    time=None):
-        #from paste.auth import auth_tkt
+                    time=None, digest_algo="md5"):
         import repoze.who._auth_tkt as auth_tkt
         ticket = auth_tkt.AuthTicket(
             'secret',
@@ -47,7 +46,8 @@
             user_data=userdata,
             time=time,
             cookie_name=cookie_name,
-            secure=secure)
+            secure=secure,
+            digest_algo=digest_algo)
         return ticket.cookie_value()
 
     def _setNowTesting(self, value):
@@ -175,6 +175,30 @@
         self.assertEqual(environ['REMOTE_USER_DATA'],'foo=123')
         self.assertEqual(environ['AUTH_TYPE'],'cookie')
 
+    def test_identify_with_alternate_hash(self):
+        plugin = self._makeOne('secret', include_ip=False, digest_algo="sha256")
+        val = self._makeTicket(userdata='foo=123', digest_algo="sha256")
+        md5_val = self._makeTicket(userdata='foo=123')
+        self.assertNotEqual(val, md5_val)
+        # md5 is 16*2 characters long, sha256 is 32*2
+        self.assertEqual(len(val), len(md5_val)+32)
+        environ = self._makeEnviron({'HTTP_COOKIE':'auth_tkt=%s' % val})
+        result = plugin.identify(environ)
+        self.assertEqual(len(result), 4)
+        self.assertEqual(result['tokens'], [''])
+        self.assertEqual(result['repoze.who.plugins.auth_tkt.userid'], 'userid')
+        self.assertEqual(result['userdata'], {'foo': '123'})
+        self.assertTrue('timestamp' in result)
+        self.assertEqual(environ['REMOTE_USER_TOKENS'], [''])
+        self.assertEqual(environ['REMOTE_USER_DATA'],'foo=123')
+        self.assertEqual(environ['AUTH_TYPE'],'cookie')
+
+    def test_identify_bad_cookie_with_alternate_hash(self):
+        plugin = self._makeOne('secret', include_ip=True, digest_algo="sha256")
+        environ = self._makeEnviron({'HTTP_COOKIE':'auth_tkt=bogus'})
+        result = plugin.identify(environ)
+        self.assertEqual(result, None)
+    
     def test_remember_creds_same(self):
         plugin = self._makeOne('secret')
         val = self._makeTicket(userid='userid', userdata='foo=123')
@@ -182,6 +206,67 @@
         result = plugin.remember(environ, {'repoze.who.userid':'userid',
                                            'userdata':{'foo': '123'}})
         self.assertEqual(result, None)
+
+    def test_remember_creds_same_alternate_hash(self):
+        plugin = self._makeOne('secret', digest_algo="sha1")
+        val = self._makeTicket(userid='userid', userdata='foo=123', digest_algo="sha1")
+        environ = self._makeEnviron({'HTTP_COOKIE':'auth_tkt=%s' % val})
+        result = plugin.remember(environ, {'repoze.who.userid':'userid',
+                                           'userdata':{'foo': '123'}})
+        self.assertEqual(result, None)
+
+    def test_remember_creds_hash_mismatch(self):
+        plugin = self._makeOne('secret', digest_algo="sha1")
+        old_val = self._makeTicket(userid='userid', userdata='foo=123', digest_algo="md5")
+        new_val = self._makeTicket(userid='userid', userdata='foo=123', digest_algo="sha1")
+        environ = self._makeEnviron({'HTTP_COOKIE':'auth_tkt=%s' % old_val})
+        result = plugin.remember(environ, {'repoze.who.userid':'userid',
+                                           'userdata':{'foo': '123'}})
+        self.assertEqual(len(result), 3)
+        self.assertEqual(result[0],
+                         ('Set-Cookie',
+                          'auth_tkt="%s"; '
+                          'Path=/' % new_val))
+        self.assertEqual(result[1],
+                         ('Set-Cookie',
+                           'auth_tkt="%s"; '
+                           'Path=/; '
+                           'Domain=localhost'
+                            % new_val))
+        self.assertEqual(result[2],
+                         ('Set-Cookie',
+                           'auth_tkt="%s"; '
+                           'Path=/; '
+                           'Domain=.localhost'
+                            % new_val))
+
+    def test_remember_creds_secure_alternate_hash(self):
+        plugin = self._makeOne('secret', secure=True, digest_algo="sha512")
+        val = self._makeTicket(userid='userid', secure=True, userdata='foo=123', digest_algo="sha512")
+        environ = self._makeEnviron()
+        result = plugin.remember(environ, {'repoze.who.userid':'userid',
+                                           'userdata':{'foo':'123'}})
+        self.assertEqual(len(result), 3)
+        self.assertEqual(result[0],
+                         ('Set-Cookie',
+                          'auth_tkt="%s"; '
+                          'Path=/; '
+                          'secure; '
+                          'HttpOnly' % val))
+        self.assertEqual(result[1],
+                         ('Set-Cookie',
+                           'auth_tkt="%s"; '
+                           'Path=/; '
+                           'Domain=localhost; '
+                           'secure; HttpOnly'
+                            % val))
+        self.assertEqual(result[2],
+                         ('Set-Cookie',
+                           'auth_tkt="%s"; '
+                           'Path=/; '
+                           'Domain=.localhost; '
+                           'secure; HttpOnly'
+                            % val))
 
     def test_remember_creds_secure(self):
         plugin = self._makeOne('secret', secure=True)
@@ -437,6 +522,22 @@
                          ('Set-Cookie',
                           'auth_tkt="%s"; Path=/' % new_val))
 
+    def test_remember_creds_reissue_alternate_hash(self):
+        import time
+        plugin = self._makeOne('secret', reissue_time=1, digest_algo="sha256")
+        old_val = self._makeTicket(userid='userid', userdata='',
+                                   time=time.time()-2, digest_algo="sha256")
+        environ = self._makeEnviron({'HTTP_COOKIE':'auth_tkt=%s' % old_val})
+        new_val = self._makeTicket(userid='userid', userdata='',
+                                   digest_algo="sha256")
+        result = plugin.remember(environ, {'repoze.who.userid':'userid',
+                                           'userdata':''})
+        self.assertEqual(type(result[0][1]), str)
+        self.assertEqual(len(result), 3)
+        self.assertEqual(result[0],
+                         ('Set-Cookie',
+                          'auth_tkt="%s"; Path=/' % new_val))
+
     def test_l10n_sane_cookie_date(self):
         from datetime import datetime
 
@@ -591,6 +692,23 @@
             userid_checker='repoze.who.plugins.auth_tkt:make_plugin')
         self.assertEqual(plugin.userid_checker, make_plugin)
 
+    def test_factory_with_alternate_hash(self):
+        from repoze.who.plugins.auth_tkt import make_plugin
+        import hashlib
+        plugin = make_plugin('secret', digest_algo="sha1")
+        self.assertEqual(plugin.digest_algo, hashlib.sha1)
+
+    def test_factory_with_alternate_hash_func(self):
+        from repoze.who.plugins.auth_tkt import make_plugin
+        import hashlib
+        plugin = make_plugin('secret', digest_algo=hashlib.sha1)
+        self.assertEqual(plugin.digest_algo, hashlib.sha1)
+
+    def test_factory_with_bogus_hash(self):
+        from repoze.who.plugins.auth_tkt import make_plugin
+        self.assertRaises(ValueError, make_plugin,
+                          secret="fiddly", digest_algo='foo23')
+
     def test_remember_max_age_unicode(self):
         from repoze.who._compat import u
         plugin = self._makeOne('secret')

--
Gitblit v1.9.3