Initial add0.4.2

1 files changed, 200 insertions, 0 deletions

diff --git a/src/saml2/mdbcache.py b/src/saml2/mdbcache.py
new file mode 100644
index 00000000..28158973
--- /dev/null
+++ b/src/saml2/mdbcache.py
@@ -0,0 +1,200 @@
+#!/usr/bin/env python
+
+__author__ = 'rolandh'
+
+from pymongo import Connection
+#import cjson
+import time
+from datetime import datetime
+
+from saml2 import time_util
+from saml2.cache import ToOld
+from saml2.time_util import TIME_FORMAT
+
+class Cache(object):
+    def __init__(self, server=None, debug=0, db=None):
+        if server:
+            connection = Connection(server)
+        else:
+            connection = Connection()
+
+        if db:
+            self._db = connection[db]
+        else:
+            self._db = connection.pysaml2
+
+        self._cache = self._db.collection
+        self.debug = debug
+
+    def delete(self, subject_id):
+        self._cache.remove({"subject_id": subject_id})
+
+    def get_identity(self, subject_id, entities=None,
+                     check_not_on_or_after=True):
+        """ Get all the identity information that has been received and
+        are still valid about the subject.
+
+        :param subject_id: The identifier of the subject
+        :param entities: The identifiers of the entities whoes assertions are
+            interesting. If the list is empty all entities are interesting.
+        :return: A 2-tuple consisting of the identity information (a
+            dictionary of attributes and values) and the list of entities
+            whoes information has timed out.
+        """
+        res = {}
+        oldees = []
+        if not entities:
+            for item in self._cache.find({"subject_id": subject_id}):
+                try:
+                    info = self._get_info(item, check_not_on_or_after)
+                except ToOld:
+                    oldees.append(item["entity_id"])
+                    continue
+
+                for key, vals in info["ava"].items():
+                    try:
+                        tmp = set(res[key]).union(set(vals))
+                        res[key] = list(tmp)
+                    except KeyError:
+                        res[key] = vals
+        else:
+            for entity_id in entities:
+                try:
+                    info = self.get(subject_id, entity_id, check_not_on_or_after)
+                except ToOld:
+                    oldees.append(entity_id)
+                    continue
+
+                for key, vals in info["ava"].items():
+                    try:
+                        tmp = set(res[key]).union(set(vals))
+                        res[key] = list(tmp)
+                    except KeyError:
+                        res[key] = vals
+
+        return res, oldees
+                
+    def _get_info(self, item, check_not_on_or_after=True):
+        """ Get session information about a subject gotten from a
+        specified IdP/AA.
+
+        :param item: Information stored
+        :return: The session information as a dictionary
+        """
+        timestamp = item["timestamp"]
+
+        if check_not_on_or_after and not time_util.not_on_or_after(timestamp):
+            raise ToOld()
+
+        try:
+            return item["info"]
+        except KeyError:
+            return None
+
+    def get(self, subject_id, entity_id, check_not_on_or_after=True):
+        res = self._cache.find_one({"subject_id": subject_id,
+                                    "entity_id": entity_id})
+        if not res:
+            return {}
+        else:
+            return self._get_info(res, check_not_on_or_after)
+
+    def set(self, subject_id, entity_id, info, timestamp=0):
+        """ Stores session information in the cache. Assumes that the subject_id
+        is unique within the context of the Service Provider.
+
+        :param subject_id: The subject identifier
+        :param entity_id: The identifier of the entity_id/receiver of an
+            assertion
+        :param info: The session info, the assertion is part of this
+        :param timestamp: A time after which the assertion is not valid.
+        """
+
+        if isinstance(timestamp, datetime) or isinstance(timestamp,
+                                                         time.struct_time):
+            timestamp = time.strftime(TIME_FORMAT, timestamp)
+
+        doc = {"subject_id": subject_id,
+               "entity_id": entity_id,
+               "info": info,
+               "timestamp": timestamp}
+
+        _ = self._cache.insert(doc)
+
+
+    def reset(self, subject_id, entity_id):
+        """ Scrap the assertions received from a IdP or an AA about a special
+        subject.
+
+        :param subject_id: The subjects identifier
+        :param entity_id: The identifier of the entity_id of the assertion
+        :return:
+        """
+        self._cache.update({"subject_id":subject_id,
+                                   "entity_id":entity_id},
+                                  {"$set": {"info":{}, "timestamp": 0}})
+
+    def entities(self, subject_id):
+        """ Returns all the entities of assertions for a subject, disregarding
+        whether the assertion still is valid or not.
+
+        :param subject_id: The identifier of the subject
+        :return: A possibly empty list of entity identifiers
+        """
+        try:
+            return [i["entity_id"] for i in self._cache.find({"subject_id":
+                    subject_id})]
+        except ValueError:
+            return []
+
+
+    def receivers(self, subject_id):
+        """ Another name for entities() just to make it more logic in the IdP
+            scenario """
+        return self.entities(subject_id)
+
+    def active(self, subject_id, entity_id):
+        """ Returns the status of assertions from a specific entity_id.
+
+        :param subject_id: The ID of the subject
+        :param entity_id: The entity ID of the entity_id of the assertion
+        :return: True or False depending on if the assertion is still
+            valid or not.
+        """
+
+        item = self._cache.find_one({"subject_id":subject_id,
+                                   "entity_id":entity_id})
+        try:
+            return time_util.not_on_or_after(item["timestamp"])
+        except ToOld:
+            return False
+
+    def subjects(self):
+        """ Return identifiers for all the subjects that are in the cache.
+
+        :return: list of subject identifiers
+        """
+
+        subj = [i["subject_id"] for i in self._cache.find()]
+
+        return list(set(subj))
+
+    def update(self, subject_id, entity_id, ava):
+        """ """
+        item = self._cache.find_one({"subject_id":subject_id,
+                                       "entity_id":entity_id})
+        info = item["info"]
+        info["ava"].update(ava)
+        self._cache.update({"subject_id":subject_id,
+                            "entity_id":entity_id},
+                            {"$set": {"info":info}})
+
+
+    def valid_to(self, subject_id, entity_id, newtime):
+        """ """
+        self._cache.update({"subject_id":subject_id,
+                            "entity_id":entity_id},
+                            {"$set": {"timestamp": newtime}})
+
+    def clear(self):
+        self._cache.remove()
\ No newline at end of file