<s> """<STR_LIT>""" <EOL> import abc <EOL> import atexit <EOL> import sys <EOL> import time <EOL> import logging <EOL> from grr . lib import access_control <EOL> from grr . lib import blob_store <EOL> from grr . lib import config_lib <EOL> from grr . lib import flags <EOL> from grr . lib import registry <EOL> from grr . lib import stats <EOL> from grr . lib import utils <EOL> flags . DEFINE_bool ( "<STR_LIT>" , False , <EOL> "<STR_LIT>" ) <EOL> DB = None <EOL> class Error ( stats . CountingExceptionMixin , Exception ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> class TimeoutError ( Exception ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> class TransactionError ( Error ) : <EOL> """<STR_LIT>""" <EOL> counter = "<STR_LIT>" <EOL> default_token = None <EOL> class MutationPool ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , token = None ) : <EOL> self . token = token <EOL> self . delete_subject_requests = [ ] <EOL> self . set_requests = [ ] <EOL> self . delete_attributes_requests = [ ] <EOL> def DeleteSubject ( self , subject ) : <EOL> self . delete_subject_requests . append ( subject ) <EOL> def MultiSet ( self , subject , values , timestamp = None , replace = True , <EOL> to_delete = None ) : <EOL> self . set_requests . append ( ( subject , values , timestamp , replace , to_delete ) ) <EOL> def Set ( self , subject , attribute , value , timestamp = None , replace = True ) : <EOL> self . MultiSet ( subject , { attribute : [ value ] } , timestamp = timestamp , <EOL> replace = replace ) <EOL> def DeleteAttributes ( self , subject , attributes , start = None , end = None ) : <EOL> self . delete_attributes_requests . append ( ( subject , attributes , start , end ) ) <EOL> def Flush ( self ) : <EOL> """<STR_LIT>""" <EOL> DB . DeleteSubjects ( self . delete_subject_requests , token = self . token , <EOL> sync = False ) <EOL> for req in self . delete_attributes_requests : <EOL> subject , attributes , start , end = req <EOL> DB . DeleteAttributes ( subject , attributes , start = start , end = end , <EOL> token = self . token , sync = False ) <EOL> for req in self . set_requests : <EOL> subject , values , timestamp , replace , to_delete = req <EOL> DB . MultiSet ( subject , values , timestamp = timestamp , replace = replace , <EOL> to_delete = to_delete , token = self . token , sync = False ) <EOL> if ( self . delete_subject_requests or <EOL> self . delete_attributes_requests or <EOL> self . set_requests ) : <EOL> DB . Flush ( ) <EOL> self . delete_subject_requests = [ ] <EOL> self . set_requests = [ ] <EOL> self . delete_attributes_requests = [ ] <EOL> def __enter__ ( self ) : <EOL> return self <EOL> def __exit__ ( self , unused_type , unused_value , unused_traceback ) : <EOL> self . Flush ( ) <EOL> def Size ( self ) : <EOL> return ( len ( self . delete_subject_requests ) + len ( self . set_requests ) + <EOL> len ( self . delete_attributes_requests ) ) <EOL> class DataStore ( object ) : <EOL> """<STR_LIT>""" <EOL> __metaclass__ = registry . MetaclassRegistry <EOL> ALL_TIMESTAMPS = "<STR_LIT>" <EOL> NEWEST_TIMESTAMP = "<STR_LIT>" <EOL> TIMESTAMPS = [ ALL_TIMESTAMPS , NEWEST_TIMESTAMP ] <EOL> mutation_pool_cls = MutationPool <EOL> flusher_thread = None <EOL> monitor_thread = None <EOL> def __init__ ( self ) : <EOL> security_manager = access_control . BasicAccessControlManager . GetPlugin ( <EOL> config_lib . CONFIG [ "<STR_LIT>" ] ) ( ) <EOL> self . security_manager = security_manager <EOL> logging . info ( "<STR_LIT>" , security_manager ) <EOL> self . flusher_thread = utils . InterruptableThread ( target = self . Flush , <EOL> sleep_time = <NUM_LIT:0.5> ) <EOL> self . flusher_thread . start ( ) <EOL> self . monitor_thread = None <EOL> def GetRequiredResolveAccess ( self , attribute_prefix ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( attribute_prefix , basestring ) : <EOL> attribute_prefix = [ utils . SmartStr ( attribute_prefix ) ] <EOL> else : <EOL> attribute_prefix = [ utils . SmartStr ( x ) for x in attribute_prefix ] <EOL> for prefix in attribute_prefix : <EOL> if not prefix : <EOL> return "<STR_LIT>" <EOL> try : <EOL> column_family , _ = prefix . split ( "<STR_LIT::>" , <NUM_LIT:1> ) <EOL> except ValueError : <EOL> raise RuntimeError ( "<STR_LIT>" <EOL> "<STR_LIT>" % prefix ) <EOL> if column_family . startswith ( "<STR_LIT:index>" ) : <EOL> return "<STR_LIT>" <EOL> return "<STR_LIT:r>" <EOL> def InitializeBlobstore ( self ) : <EOL> blobstore_name = config_lib . CONFIG . Get ( "<STR_LIT>" ) <EOL> try : <EOL> cls = blob_store . Blobstore . GetPlugin ( blobstore_name ) <EOL> except KeyError : <EOL> raise RuntimeError ( "<STR_LIT>" % blobstore_name ) <EOL> self . blobstore = cls ( ) <EOL> def InitializeMonitorThread ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . monitor_thread : <EOL> return <EOL> self . monitor_thread = utils . InterruptableThread ( target = self . _RegisterSize , <EOL> sleep_time = <NUM_LIT> ) <EOL> self . monitor_thread . start ( ) <EOL> def _RegisterSize ( self ) : <EOL> """<STR_LIT>""" <EOL> stats . STATS . SetGaugeValue ( "<STR_LIT>" , self . Size ( ) ) <EOL> def Initialize ( self ) : <EOL> """<STR_LIT>""" <EOL> self . InitializeBlobstore ( ) <EOL> @ abc . abstractmethod <EOL> def DeleteSubject ( self , subject , sync = False , token = None ) : <EOL> """<STR_LIT>""" <EOL> def DeleteSubjects ( self , subjects , sync = False , token = None ) : <EOL> """<STR_LIT>""" <EOL> for subject in subjects : <EOL> self . DeleteSubject ( subject , sync = sync , token = token ) <EOL> def Set ( self , subject , attribute , value , timestamp = None , token = None , <EOL> replace = True , sync = True ) : <EOL> """<STR_LIT>""" <EOL> self . MultiSet ( subject , { attribute : [ value ] } , timestamp = timestamp , <EOL> token = token , replace = replace , sync = sync ) <EOL> def RetryWrapper ( self , subject , callback , retrywrap_timeout = <NUM_LIT:1> , token = None , <EOL> retrywrap_max_timeout = <NUM_LIT:10> , ** kw ) : <EOL> """<STR_LIT>""" <EOL> def Retry ( ) : <EOL> """<STR_LIT>""" <EOL> transaction = self . Transaction ( subject , token = token ) <EOL> try : <EOL> result = callback ( transaction , ** kw ) <EOL> finally : <EOL> transaction . Commit ( ) <EOL> return result <EOL> timeout = <NUM_LIT:0> <EOL> while timeout < retrywrap_max_timeout : <EOL> try : <EOL> result = Retry ( ) <EOL> if timeout > <NUM_LIT:1> : <EOL> logging . debug ( "<STR_LIT>" , timeout ) <EOL> return result <EOL> except TransactionError : <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" ) <EOL> time . sleep ( retrywrap_timeout ) <EOL> timeout += retrywrap_timeout <EOL> raise TransactionError ( "<STR_LIT>" ) <EOL> @ abc . abstractmethod <EOL> def Transaction ( self , subject , lease_time = None , token = None ) : <EOL> """<STR_LIT>""" <EOL> @ abc . abstractmethod <EOL> def MultiSet ( self , subject , values , timestamp = None , replace = True , sync = True , <EOL> to_delete = None , token = None ) : <EOL> """<STR_LIT>""" <EOL> def MultiDeleteAttributes ( self , subjects , attributes , start = None , end = None , <EOL> sync = True , token = None ) : <EOL> """<STR_LIT>""" <EOL> for subject in subjects : <EOL> self . DeleteAttributes ( subject , attributes , start = start , end = end , <EOL> sync = sync , token = token ) <EOL> @ abc . abstractmethod <EOL> def DeleteAttributes ( self , subject , attributes , start = None , end = None , <EOL> sync = True , token = None ) : <EOL> """<STR_LIT>""" <EOL> def Resolve ( self , subject , attribute , token = None ) : <EOL> """<STR_LIT>""" <EOL> for _ , value , timestamp in self . ResolveMulti ( <EOL> subject , [ attribute ] , token = token , timestamp = self . NEWEST_TIMESTAMP ) : <EOL> return value , timestamp <EOL> return ( None , <NUM_LIT:0> ) <EOL> @ abc . abstractmethod <EOL> def MultiResolvePrefix ( self , subjects , attribute_prefix , timestamp = None , <EOL> limit = None , token = None ) : <EOL> """<STR_LIT>""" <EOL> def ResolvePrefix ( self , subject , attribute_prefix , timestamp = None , <EOL> limit = None , token = None ) : <EOL> """<STR_LIT>""" <EOL> for _ , values in self . MultiResolvePrefix ( <EOL> [ subject ] , attribute_prefix , timestamp = timestamp , token = token , <EOL> limit = limit ) : <EOL> values . sort ( key = lambda a : a [ <NUM_LIT:0> ] ) <EOL> return values <EOL> return [ ] <EOL> def ResolveMulti ( self , subject , attributes , timestamp = None , limit = None , <EOL> token = None ) : <EOL> """<STR_LIT>""" <EOL> def ResolveRow ( self , subject , ** kw ) : <EOL> return self . ResolvePrefix ( subject , "<STR_LIT>" , ** kw ) <EOL> def Flush ( self ) : <EOL> """<STR_LIT>""" <EOL> def Size ( self ) : <EOL> """<STR_LIT>""" <EOL> return - <NUM_LIT:1> <EOL> def __del__ ( self ) : <EOL> if self . flusher_thread : <EOL> self . flusher_thread . Stop ( ) <EOL> if self . monitor_thread : <EOL> self . monitor_thread . Stop ( ) <EOL> try : <EOL> self . Flush ( ) <EOL> except Exception : <EOL> pass <EOL> @ abc . abstractmethod <EOL> def ScanAttributes ( self , subject_prefix , attributes , <EOL> after_urn = None , max_records = None , token = None , <EOL> relaxed_order = False ) : <EOL> """<STR_LIT>""" <EOL> def ScanAttribute ( self , subject_prefix , attribute , <EOL> after_urn = None , max_records = None , token = None , <EOL> relaxed_order = False ) : <EOL> for s , r in self . ScanAttributes ( subject_prefix , [ attribute ] , <EOL> after_urn = after_urn , <EOL> max_records = max_records , <EOL> token = token , <EOL> relaxed_order = relaxed_order ) : <EOL> ts , v = r [ attribute ] <EOL> yield ( s , ts , v ) <EOL> def ReadBlob ( self , identifier , token = None ) : <EOL> return self . ReadBlobs ( [ identifier ] , token = token ) . values ( ) [ <NUM_LIT:0> ] <EOL> def ReadBlobs ( self , identifiers , token = None ) : <EOL> return self . blobstore . ReadBlobs ( identifiers , token = token ) <EOL> def StoreBlob ( self , content , token = None ) : <EOL> return self . blobstore . StoreBlob ( content , token = token ) <EOL> def StoreBlobs ( self , contents , token = None ) : <EOL> return self . blobstore . StoreBlobs ( contents , token = token ) <EOL> def BlobExists ( self , identifier , token = None ) : <EOL> return self . BlobsExist ( [ identifier ] , token = token ) . values ( ) [ <NUM_LIT:0> ] <EOL> def BlobsExist ( self , identifiers , token = None ) : <EOL> return self . blobstore . BlobsExist ( identifiers , token = token ) <EOL> def GetMutationPool ( self , token = None ) : <EOL> return self . mutation_pool_cls ( token = token ) <EOL> class Transaction ( object ) : <EOL> """<STR_LIT>""" <EOL> __metaclass__ = registry . MetaclassRegistry <EOL> @ abc . abstractmethod <EOL> def __init__ ( self , table , subject , lease_time = None , token = None ) : <EOL> """<STR_LIT>""" <EOL> @ abc . abstractmethod <EOL> def DeleteAttribute ( self , attribute ) : <EOL> """<STR_LIT>""" <EOL> @ abc . abstractmethod <EOL> def Resolve ( self , attribute ) : <EOL> """<STR_LIT>""" <EOL> @ abc . abstractmethod <EOL> def ResolvePrefix ( self , attribute_prefix , timestamp = None ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> def UpdateLease ( self , duration ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError <EOL> def CheckLease ( self ) : <EOL> """<STR_LIT>""" <EOL> return True <EOL> @ abc . abstractmethod <EOL> def Set ( self , attribute , value , timestamp = None , replace = True ) : <EOL> """<STR_LIT>""" <EOL> @ abc . abstractmethod <EOL> def Commit ( self ) : <EOL> """<STR_LIT>""" <EOL> @ abc . abstractmethod <EOL> def Abort ( self ) : <EOL> """<STR_LIT>""" <EOL> class CommonTransaction ( Transaction ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , table , subject , lease_time = None , token = None ) : <EOL> super ( CommonTransaction , self ) . __init__ ( table , subject , <EOL> lease_time = lease_time , token = token ) <EOL> self . to_set = { } <EOL> self . to_delete = set ( ) <EOL> self . subject = subject <EOL> self . store = table <EOL> self . token = token <EOL> self . expires = None <EOL> def CheckLease ( self ) : <EOL> if not self . expires : <EOL> return <NUM_LIT:0> <EOL> return max ( <NUM_LIT:0> , self . expires - time . time ( ) ) <EOL> def DeleteAttribute ( self , attribute ) : <EOL> self . to_delete . add ( attribute ) <EOL> def ResolvePrefix ( self , prefix , timestamp = None ) : <EOL> if isinstance ( timestamp , ( list , tuple ) ) : <EOL> start , end = timestamp <EOL> elif isinstance ( timestamp , int ) : <EOL> start = timestamp <EOL> end = timestamp <EOL> elif timestamp == DataStore . ALL_TIMESTAMPS or timestamp is None : <EOL> start , end = <NUM_LIT:0> , ( <NUM_LIT:2> ** <NUM_LIT> ) - <NUM_LIT:1> <EOL> timestamp = ( start , end ) <EOL> elif timestamp == DataStore . NEWEST_TIMESTAMP : <EOL> start , end = <NUM_LIT:0> , ( <NUM_LIT:2> ** <NUM_LIT> ) - <NUM_LIT:1> <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % <EOL> utils . SmartStr ( timestamp ) ) <EOL> start = int ( start ) <EOL> end = int ( end ) <EOL> results = [ ] <EOL> if self . to_set : <EOL> for attribute , values in self . to_set . items ( ) : <EOL> if utils . SmartStr ( attribute ) . startswith ( prefix ) : <EOL> results . extend ( [ ( attribute , value , ts ) for value , ts in values <EOL> if start <= ts <= end ] ) <EOL> ds_results = self . store . ResolvePrefix ( self . subject , prefix , <EOL> timestamp = timestamp , <EOL> token = self . token ) <EOL> if self . to_delete : <EOL> for val in ds_results : <EOL> attribute , value , ts = val <EOL> if attribute not in self . to_delete : <EOL> results . append ( val ) <EOL> else : <EOL> results . extend ( ds_results ) <EOL> if timestamp == DataStore . NEWEST_TIMESTAMP : <EOL> newest_results = { } <EOL> for attribute , value , ts in results : <EOL> current = newest_results . get ( attribute , None ) <EOL> if current : <EOL> _ , _ , current_ts = current <EOL> if ts > current_ts : <EOL> newest_results [ attribute ] = ( attribute , value , ts ) <EOL> else : <EOL> newest_results [ attribute ] = ( attribute , value , ts ) <EOL> results = newest_results . values ( ) <EOL> return sorted ( results , key = lambda ( a , val , ts ) : ( a , ts , val ) ) <EOL> def Set ( self , attribute , value , timestamp = None , replace = True ) : <EOL> if replace : <EOL> self . to_delete . add ( attribute ) <EOL> if timestamp is None : <EOL> timestamp = int ( time . time ( ) * <NUM_LIT> ) <EOL> self . to_set . setdefault ( attribute , [ ] ) . append ( ( value , int ( timestamp ) ) ) <EOL> def Resolve ( self , attribute ) : <EOL> if attribute in self . to_set : <EOL> return max ( self . to_set [ attribute ] , key = lambda vt : vt [ <NUM_LIT:1> ] ) <EOL> if attribute in self . to_delete : <EOL> return ( None , <NUM_LIT:0> ) <EOL> return self . store . Resolve ( self . subject , attribute , token = self . token ) <EOL> def Commit ( self ) : <EOL> if not self . CheckLease ( ) : <EOL> raise TransactionError ( "<STR_LIT>" ) <EOL> self . store . DeleteAttributes ( self . subject , self . to_delete , sync = True , <EOL> token = self . token ) <EOL> self . store . MultiSet ( self . subject , self . to_set , token = self . token ) <EOL> self . to_set = { } <EOL> self . to_delete = set ( ) <EOL> def __del__ ( self ) : <EOL> try : <EOL> self . Abort ( ) <EOL> except Exception : <EOL> pass <EOL> class ResultSet ( object ) : <EOL> """<STR_LIT>""" <EOL> total_count = <NUM_LIT:0> <EOL> def __init__ ( self , results = None ) : <EOL> if results is None : <EOL> results = [ ] <EOL> self . results = results <EOL> def __iter__ ( self ) : <EOL> return iter ( self . results ) <EOL> def __getitem__ ( self , item ) : <EOL> return self . results [ item ] <EOL> def __len__ ( self ) : <EOL> return len ( self . results ) <EOL> def __iadd__ ( self , other ) : <EOL> self . results = list ( self . results ) + list ( other ) <EOL> return self <EOL> def Append ( self , item ) : <EOL> self . results . append ( item ) <EOL> class DataStoreInit ( registry . InitHook ) : <EOL> """<STR_LIT>""" <EOL> pre = [ "<STR_LIT>" ] <EOL> def Run ( self ) : <EOL> """<STR_LIT>""" <EOL> global DB <EOL> if flags . FLAGS . list_storage : <EOL> for name , cls in DataStore . classes . items ( ) : <EOL> print "<STR_LIT>" % ( name , cls . __doc__ ) <EOL> sys . exit ( <NUM_LIT:0> ) <EOL> try : <EOL> cls = DataStore . GetPlugin ( config_lib . CONFIG [ "<STR_LIT>" ] ) <EOL> except KeyError : <EOL> raise RuntimeError ( "<STR_LIT>" % <EOL> config_lib . CONFIG [ "<STR_LIT>" ] ) <EOL> DB = cls ( ) <EOL> DB . Initialize ( ) <EOL> atexit . register ( DB . Flush ) <EOL> monitor_port = config_lib . CONFIG [ "<STR_LIT>" ] <EOL> if monitor_port != <NUM_LIT:0> : <EOL> stats . STATS . RegisterGaugeMetric ( "<STR_LIT>" , int , <EOL> docstring = "<STR_LIT>" , <EOL> units = "<STR_LIT>" ) <EOL> DB . InitializeMonitorThread ( ) <EOL> def RunOnce ( self ) : <EOL> """<STR_LIT>""" <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) </s>
