added dbapi2 (sql) backed noncestore! yeah :-)

1 files changed, 134 insertions, 0 deletions

diff --git a/wsgitools/digest.py b/wsgitools/digest.py
index 52cdcee..d0039e0 100755
--- a/wsgitools/digest.py
+++ b/wsgitools/digest.py
@@ -293,6 +293,140 @@ class MemoryNonceStore(NonceStoreBase):
             self.nonces[lower] = (nt, nv, uses+1)
         return True
 
+__all__.append("LazyDBAPI2Opener")
+class LazyDBAPI2Opener:
+    """
+    Connects to database on first request. Otherwise it behaves like a dbapi2
+    connection.
+    """
+    def __init__(self, function, *args, **kwargs):
+        """
+        The database will be connected on the first method call. This is done
+        by calling the given function with the remaining parameters.
+        @param function: is the function that connects to the database
+        """
+        self._function = function
+        self._args = args
+        self._kwargs = kwargs
+        self._dbhandle = None
+    def _getdbhandle(self):
+        """Returns an open database connection. Open if necessary."""
+        if self._dbhandle is None:
+            self._dbhandle = self._function(*self._args, **self._kwargs)
+            self._function = self._args = self._kwargs = None
+        return self._dbhandle
+    def cursor(self):
+        """dbapi2"""
+        return self._getdbhandle().cursor()
+    def commit(self):
+        """dbapi2"""
+        return self._getdbhandle().commit()
+    def rollback(self):
+        """dbapi2"""
+        return self._getdbhandle().rollback()
+    def close(self):
+        """dbapi2"""
+        return self._getdbhandle().close()
+
+__all__.append("DBAPI2NonceStore")
+class DBAPI2NonceStore(NonceStoreBase):
+    """
+    A dbapi2-backed nonce store implementation suitable for usage with forking
+    wsgi servers such as scgi.forkpool.
+    """
+    def __init__(self, dbhandle, maxage=300, maxuses=5, table="nonces"):
+        """
+        @type filename: str
+        @param filename: the path to the bsddb file that is used as backingstore
+        @type maxage: int
+        @param maxage: is the number of seconds a nonce may be valid. Choosing a
+               large value may result in more memory usage whereas a smaller
+               value results in more requests. Defaults to 5 minutes.
+        @type maxuses: int
+        @param maxuses: is the number of times a nonce may be used (with
+                different nc values). A value of 1 makes nonces usable exactly
+                once resulting in more requests. Defaults to 5.
+        """
+        NonceStoreBase.__init__(self)
+        self.dbhandle = dbhandle
+        self.maxage = maxage
+        self.maxuses = maxuses
+        self.table = table
+        self.server_secret = gen_rand_str()
+
+    def _cleanup(self, cur):
+        """internal methods cleaning list of valid nonces"""
+        old = format_time(time.time() - self.maxage)
+        cur.execute("DELETE FROM %s WHERE key < '%s:';" % (self.table, old))
+
+    def newnonce(self, ident=None):
+        """
+        Generates a new nonce string.
+        @rtype: str
+        """
+        nonce_time = format_time(time.time())
+        nonce_value = gen_rand_str()
+        dbkey = "%s:%s" % (nonce_time, nonce_value)
+        cur = self.dbhandle.cursor()
+        self._cleanup(cur) # avoid growing database
+        cur.execute("INSERT INTO %s VALUES ('%s', '1');" % (self.table, dbkey))
+        self.dbhandle.commit()
+        token = "%s:%s" % (dbkey, self.server_secret)
+        if ident is not None:
+            token = "%s:%s" % (token, ident)
+        token = md5(token).hexdigest()
+        return "%s:%s:%s" % (nonce_time, nonce_value, token)
+
+    def checknonce(self, nonce, count=1, ident=None):
+        """
+        Do a check for whether the provided string is a nonce and increase usage
+        count on returning True.
+        @type nonce: str
+        @type count: int
+        @rtype: bool
+        """
+        try:
+            nonce_time, nonce_value, nonce_hash = nonce.split(':')
+        except ValueError:
+            return False
+        if not nonce_time.isalnum() or not nonce_value.replace("+", ""). \
+           replace("/", "").replace("=", "").isalnum():
+            return False
+        token = "%s:%s:%s" % (nonce_time, nonce_value, self.server_secret)
+        if ident is not None:
+            token = "%s:%s" % (token, ident)
+        token = md5(token).hexdigest()
+        if token != nonce_hash:
+            return False
+
+        if nonce_time < format_time(time.time() - self.maxage):
+            return False
+
+        cur = self.dbhandle.cursor()
+        #self._cleanup(cur) # avoid growing database
+
+        dbkey = "%s:%s" % (nonce_time, nonce_value)
+        cur.execute("SELECT value FROM %s WHERE key = '%s';" %
+                    (self.table, dbkey))
+        uses = cur.fetchone()
+        if uses is None:
+            self.dbhandle.commit()
+            return False
+        uses = int(uses[0])
+        if count != uses:
+            cur.execute("DELETE FROM %s WHERE key = '%s';" %
+                        (self.table, dbkey))
+            self.dbhandle.commit()
+            return False
+        if uses >= self.maxuses:
+            cur.execute("DELETE FROM %s WHERE key = '%s';" %
+                        (self.table, dbkey))
+        else:
+            cur.execute("UPDATE %s SET value = '%d' WHERE key = '%s';" %
+                        (self.table, uses + 1, dbkey))
+        self.dbhandle.commit()
+        return True
+
 __all__.append("AuthDigestMiddleware")
 class AuthDigestMiddleware:
     """Middleware partly implementing RFC2617. (md5-sess was omited)"""