Memcache Storage for Flask-OpenId

memcache_storage_flask_openid

from openid.store.interface import OpenIDStore
from openid.association import Association
from pymemcache.client.base import Client

# use pymemcache as the memcache client lib.
# memcache client should have way to serializer and deserializer data to and from memcache store. So that we can put python dict in memcache
# pymemcache has this support if in client initialization you pass serializer and deserializer functions.

def json_serializer(key, value):
    if type(value) == str:
        return value, 1
    return json.dumps(value), 2


def json_deserializer(key, value, flags):
    if flags == 1:
        return value
    if flags == 2:
        return json.loads(value)
    raise Exception("Unknown serialization format")

memcache = Client(('localhost', 11211), serializer=json_serializer, deserializer=json_deserializer)

class MemcacheStore(OpenIDStore):
    def association_key(self, server_url, handle=None):
        """
        Memcache key is prefixed with 'openid_association_' string
        """
        return 'openid_association_' + server_url + '_' + handle

    def association_server_key(self, server_url):
        """
        Memcache key is prefixed with 'openid_association_' string.
        """
        return 'openid_association_server_' + server_url

    def storeAssociation(self, server_url, association):
        # create memcache key for association itself
        # and list of associations for this server
        association_key = self.association_key(server_url, association.handle)
        server_key = self.association_server_key(server_url)

        # get list of associations
        server_associations = memcache.get(server_key)

        # if no such list, initialize it with empty list
        if not server_associations:
            server_associations = {}

        # and store given association key in it
        server_associations[association.issued] = association_key

        # save association's keys list
        memcache.set(server_key, server_associations)

        # save association itself
        memcache.set(association_key, association.serialize(), 30)

    def getAssociation(self, server_url, handle=None):
        # simple case: handle given
        if handle:
            assoc_key = self.association_key(server_url, handle)
            association_data = memcache.get(assoc_key)

            if association_data:
                assoc = Association.deserialize(association_data)
                return assoc
            else:
                return None

        # no handle given, working with list
        # create key for list of associations
        server_key = self.association_server_key(server_url)

        # get list of associations
        server_associations = memcache.get(server_key)

        if not server_associations:
            return None
        keys = server_associations.keys()
        keys.sort()
        last_key = server_associations[keys[-1]]

        # get association, return null if failed
        association_data = memcache.get(last_key)

        if association_data:
            assoc = Association.deserialize(association_data)
            return assoc
        else:
            return None

    def removeAssociation(self, server_url, handle):
        server_key = self.association_server_key(server_url)
        association_key = self.association_key(server_url, handle)

        server_associations = memcache.get(server_key)

        if not server_associations:
            return False
        server_associations = {v: k for k, v in server_associations.items()}

        # ensure that given association key exists in list
        if not association_key in server_associations:
            return False

        # remove given association key from list
        server_associations.pop(association_key)
        server_associations = {v: k for k, v in server_associations.items()}

        # save updated list
        memcache.set(server_key, server_associations)

        # delete association
        return memcache.delete(association_key)

    def useNonce(self, server_url, timestamp, salt):
        nonce = ['openid_nonce', server_url, str(timestamp), salt]
        nonce_key = '_'.join(nonce)
        return memcache.add(nonce_key, '-', 60)

    def cleanupNonces(self):
        None

    def cleanupAssociations(self):
        None

    def cleanup(self):
        None