summaryrefslogtreecommitdiff
path: root/src/saml2/mdbcache.py
blob: 6da1b498e0d6d1f1d23de5d3e1e9733dc1576b70 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
#!/usr/bin/env python
import logging
from pymongo.mongo_client import MongoClient

__author__ = 'rolandh'

import time
from datetime import datetime

from saml2 import time_util
from saml2.cache import TooOld
from saml2.time_util import TIME_FORMAT

logger = logging.getLogger(__name__)


class Cache(object):
    def __init__(self, server=None, debug=0, db=None):
        if server:
            connection = MongoClient(server)
        else:
            connection = MongoClient()

        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.delete_many({'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 TooOld:
                    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 TooOld:
                    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 TooOld()

        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_one(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 TooOld:
            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.delete_many({})