<s> import logging <EOL> __author__ = '<STR_LIT>' <EOL> from pymongo import Connection <EOL> import time <EOL> from datetime import datetime <EOL> from saml2 import time_util <EOL> from saml2 . cache import ToOld <EOL> from saml2 . time_util import TIME_FORMAT <EOL> logger = logging . getLogger ( __name__ ) <EOL> class Cache ( object ) : <EOL> def __init__ ( self , server = None , debug = <NUM_LIT:0> , db = None ) : <EOL> if server : <EOL> connection = Connection ( server ) <EOL> else : <EOL> connection = Connection ( ) <EOL> if db : <EOL> self . _db = connection [ db ] <EOL> else : <EOL> self . _db = connection . pysaml2 <EOL> self . _cache = self . _db . collection <EOL> self . debug = debug <EOL> def delete ( self , subject_id ) : <EOL> self . _cache . remove ( { "<STR_LIT>" : subject_id } ) <EOL> def get_identity ( self , subject_id , entities = None , <EOL> check_not_on_or_after = True ) : <EOL> """<STR_LIT>""" <EOL> res = { } <EOL> oldees = [ ] <EOL> if not entities : <EOL> for item in self . _cache . find ( { "<STR_LIT>" : subject_id } ) : <EOL> try : <EOL> info = self . _get_info ( item , check_not_on_or_after ) <EOL> except ToOld : <EOL> oldees . append ( item [ "<STR_LIT>" ] ) <EOL> continue <EOL> for key , vals in info [ "<STR_LIT>" ] . items ( ) : <EOL> try : <EOL> tmp = set ( res [ key ] ) . union ( set ( vals ) ) <EOL> res [ key ] = list ( tmp ) <EOL> except KeyError : <EOL> res [ key ] = vals <EOL> else : <EOL> for entity_id in entities : <EOL> try : <EOL> info = self . get ( subject_id , entity_id , <EOL> check_not_on_or_after ) <EOL> except ToOld : <EOL> oldees . append ( entity_id ) <EOL> continue <EOL> for key , vals in info [ "<STR_LIT>" ] . items ( ) : <EOL> try : <EOL> tmp = set ( res [ key ] ) . union ( set ( vals ) ) <EOL> res [ key ] = list ( tmp ) <EOL> except KeyError : <EOL> res [ key ] = vals <EOL> return res , oldees <EOL> def _get_info ( self , item , check_not_on_or_after = True ) : <EOL> """<STR_LIT>""" <EOL> timestamp = item [ "<STR_LIT>" ] <EOL> if check_not_on_or_after and not time_util . not_on_or_after ( timestamp ) : <EOL> raise ToOld ( ) <EOL> try : <EOL> return item [ "<STR_LIT:info>" ] <EOL> except KeyError : <EOL> return None <EOL> def get ( self , subject_id , entity_id , check_not_on_or_after = True ) : <EOL> res = self . _cache . find_one ( { "<STR_LIT>" : subject_id , <EOL> "<STR_LIT>" : entity_id } ) <EOL> if not res : <EOL> return { } <EOL> else : <EOL> return self . _get_info ( res , check_not_on_or_after ) <EOL> def set ( self , subject_id , entity_id , info , timestamp = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( timestamp , datetime ) or isinstance ( timestamp , <EOL> time . struct_time ) : <EOL> timestamp = time . strftime ( TIME_FORMAT , timestamp ) <EOL> doc = { "<STR_LIT>" : subject_id , <EOL> "<STR_LIT>" : entity_id , <EOL> "<STR_LIT:info>" : info , <EOL> "<STR_LIT>" : timestamp } <EOL> _ = self . _cache . insert ( doc ) <EOL> def reset ( self , subject_id , entity_id ) : <EOL> """<STR_LIT>""" <EOL> self . _cache . update ( { "<STR_LIT>" : subject_id , "<STR_LIT>" : entity_id } , <EOL> { "<STR_LIT>" : { "<STR_LIT:info>" : { } , "<STR_LIT>" : <NUM_LIT:0> } } ) <EOL> def entities ( self , subject_id ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> return [ i [ "<STR_LIT>" ] for i in self . _cache . find ( { "<STR_LIT>" : <EOL> subject_id } ) ] <EOL> except ValueError : <EOL> return [ ] <EOL> def receivers ( self , subject_id ) : <EOL> """<STR_LIT>""" <EOL> return self . entities ( subject_id ) <EOL> def active ( self , subject_id , entity_id ) : <EOL> """<STR_LIT>""" <EOL> item = self . _cache . find_one ( { "<STR_LIT>" : subject_id , <EOL> "<STR_LIT>" : entity_id } ) <EOL> try : <EOL> return time_util . not_on_or_after ( item [ "<STR_LIT>" ] ) <EOL> except ToOld : <EOL> return False <EOL> def subjects ( self ) : <EOL> """<STR_LIT>""" <EOL> subj = [ i [ "<STR_LIT>" ] for i in self . _cache . find ( ) ] <EOL> return list ( set ( subj ) ) <EOL> def update ( self , subject_id , entity_id , ava ) : <EOL> """<STR_LIT:U+0020>""" <EOL> item = self . _cache . find_one ( { "<STR_LIT>" : subject_id , <EOL> "<STR_LIT>" : entity_id } ) <EOL> info = item [ "<STR_LIT:info>" ] <EOL> info [ "<STR_LIT>" ] . update ( ava ) <EOL> self . _cache . update ( { "<STR_LIT>" : subject_id , "<STR_LIT>" : entity_id } , <EOL> { "<STR_LIT>" : { "<STR_LIT:info>" : info } } ) <EOL> def valid_to ( self , subject_id , entity_id , newtime ) : <EOL> """<STR_LIT:U+0020>""" <EOL> self . _cache . update ( { "<STR_LIT>" : subject_id , "<STR_LIT>" : entity_id } , <EOL> { "<STR_LIT>" : { "<STR_LIT>" : newtime } } ) <EOL> def clear ( self ) : <EOL> self . _cache . remove ( ) </s>
