<s> """<STR_LIT>""" <EOL> import cStringIO <EOL> import itertools <EOL> import struct <EOL> import logging <EOL> from grr . lib import aff4 <EOL> from grr . lib import config_lib <EOL> from grr . lib import data_store <EOL> from grr . lib import rdfvalue <EOL> from grr . lib import registry <EOL> from grr . lib import stats <EOL> from grr . lib import utils <EOL> from grr . lib . rdfvalues import anomaly as rdf_anomaly <EOL> from grr . lib . rdfvalues import client as rdf_client <EOL> from grr . lib . rdfvalues import crypto as rdf_crypto <EOL> from grr . lib . rdfvalues import flows as rdf_flows <EOL> from grr . lib . rdfvalues import protodict as rdf_protodict <EOL> from grr . lib . rdfvalues import structs as rdf_structs <EOL> from grr . proto import jobs_pb2 <EOL> class ComponentObject ( aff4 . AFF4Object ) : <EOL> """<STR_LIT>""" <EOL> class SchemaCls ( aff4 . AFF4Object . SchemaCls ) : <EOL> COMPONENT = aff4 . Attribute ( <EOL> "<STR_LIT>" , rdf_client . ClientComponentSummary , <EOL> "<STR_LIT>" ) <EOL> class AFF4CollectionView ( rdf_protodict . RDFValueArray ) : <EOL> """<STR_LIT>""" <EOL> class RDFValueCollectionView ( rdf_protodict . RDFValueArray ) : <EOL> """<STR_LIT>""" <EOL> class RDFValueCollection ( aff4 . AFF4Object ) : <EOL> """<STR_LIT>""" <EOL> _rdf_type = None <EOL> _behaviours = set ( ) <EOL> size = <NUM_LIT:0> <EOL> fd = None <EOL> class SchemaCls ( aff4 . AFF4Object . SchemaCls ) : <EOL> SIZE = aff4 . AFF4Stream . SchemaCls . SIZE <EOL> DESCRIPTION = aff4 . Attribute ( "<STR_LIT>" , rdfvalue . RDFString , <EOL> "<STR_LIT>" , "<STR_LIT:description>" ) <EOL> VIEW = aff4 . Attribute ( "<STR_LIT>" , RDFValueCollectionView , <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" , default = "<STR_LIT>" ) <EOL> def Initialize ( self ) : <EOL> """<STR_LIT>""" <EOL> self . stream_dirty = False <EOL> try : <EOL> self . fd = aff4 . FACTORY . Open ( self . urn . Add ( "<STR_LIT>" ) , <EOL> aff4_type = "<STR_LIT>" , <EOL> mode = self . mode , <EOL> token = self . token ) <EOL> self . size = int ( self . Get ( self . Schema . SIZE ) ) <EOL> return <EOL> except IOError : <EOL> pass <EOL> try : <EOL> self . fd = aff4 . FACTORY . Open ( self . urn . Add ( "<STR_LIT>" ) , <EOL> aff4_type = "<STR_LIT>" , mode = self . mode , <EOL> token = self . token ) <EOL> self . size = int ( self . Get ( self . Schema . SIZE ) ) <EOL> return <EOL> except IOError : <EOL> pass <EOL> self . fd = aff4 . FACTORY . Create ( self . urn . Add ( "<STR_LIT>" ) , <EOL> "<STR_LIT>" , <EOL> mode = self . mode , token = self . token ) <EOL> self . fd . seek ( <NUM_LIT:0> , <NUM_LIT:2> ) <EOL> self . size = <NUM_LIT:0> <EOL> def SetChunksize ( self , chunk_size ) : <EOL> if self . fd . size != <NUM_LIT:0> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> self . fd . SetChunksize ( chunk_size ) <EOL> def Flush ( self , sync = False ) : <EOL> if self . stream_dirty : <EOL> self . Set ( self . Schema . SIZE ( self . size ) ) <EOL> self . fd . Flush ( sync = sync ) <EOL> super ( RDFValueCollection , self ) . Flush ( sync = sync ) <EOL> def Close ( self , sync = False ) : <EOL> if self . locked : <EOL> sync = True <EOL> self . Flush ( sync = sync ) <EOL> super ( RDFValueCollection , self ) . Close ( sync = sync ) <EOL> def Add ( self , rdf_value = None , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if rdf_value is None : <EOL> if self . _rdf_type : <EOL> rdf_value = self . _rdf_type ( ** kwargs ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if self . _rdf_type and not isinstance ( rdf_value , self . _rdf_type ) : <EOL> raise ValueError ( "<STR_LIT>" % <EOL> self . _rdf_type . __name__ ) <EOL> if not rdf_value . age : <EOL> rdf_value . age . Now ( ) <EOL> data = rdf_protodict . EmbeddedRDFValue ( payload = rdf_value ) . SerializeToString ( ) <EOL> self . fd . Seek ( <NUM_LIT:0> , <NUM_LIT:2> ) <EOL> self . fd . Write ( struct . pack ( "<STR_LIT>" , len ( data ) ) ) <EOL> self . fd . Write ( data ) <EOL> self . stream_dirty = True <EOL> self . size += <NUM_LIT:1> <EOL> def AddAsMessage ( self , rdfvalue_in , source ) : <EOL> """<STR_LIT>""" <EOL> self . Add ( rdf_flows . GrrMessage ( payload = rdfvalue_in , source = source ) ) <EOL> def AddAll ( self , rdf_values , callback = None ) : <EOL> """<STR_LIT>""" <EOL> for rdf_value in rdf_values : <EOL> if rdf_value is None : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if self . _rdf_type and not isinstance ( rdf_value , self . _rdf_type ) : <EOL> raise ValueError ( "<STR_LIT>" % <EOL> self . _rdf_type . __name__ ) <EOL> if not rdf_value . age : <EOL> rdf_value . age . Now ( ) <EOL> buf = cStringIO . StringIO ( ) <EOL> for index , rdf_value in enumerate ( rdf_values ) : <EOL> data = rdf_protodict . EmbeddedRDFValue ( <EOL> payload = rdf_value ) . SerializeToString ( ) <EOL> buf . write ( struct . pack ( "<STR_LIT>" , len ( data ) ) ) <EOL> buf . write ( data ) <EOL> self . size += <NUM_LIT:1> <EOL> if callback : <EOL> callback ( index , rdf_value ) <EOL> self . fd . Seek ( <NUM_LIT:0> , <NUM_LIT:2> ) <EOL> self . fd . Write ( buf . getvalue ( ) ) <EOL> self . stream_dirty = True <EOL> def __len__ ( self ) : <EOL> return self . size <EOL> def __nonzero__ ( self ) : <EOL> return self . size != <NUM_LIT:0> <EOL> def __iter__ ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . GenerateItems ( ) <EOL> @ property <EOL> def deprecated_current_offset ( self ) : <EOL> return self . fd . Tell ( ) <EOL> def _GenerateItems ( self , byte_offset = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> if not self . fd : <EOL> return <EOL> if self . mode == "<STR_LIT:w>" : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> self . fd . seek ( byte_offset ) <EOL> count = <NUM_LIT:0> <EOL> while True : <EOL> offset = self . fd . Tell ( ) <EOL> try : <EOL> length = struct . unpack ( "<STR_LIT>" , self . fd . Read ( <NUM_LIT:4> ) ) [ <NUM_LIT:0> ] <EOL> serialized_event = self . fd . Read ( length ) <EOL> except struct . error : <EOL> break <EOL> result = rdf_protodict . EmbeddedRDFValue ( serialized_event ) <EOL> payload = result . payload <EOL> if payload is not None : <EOL> payload . id = count <EOL> payload . collection_offset = offset <EOL> yield payload <EOL> else : <EOL> logging . warning ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" , self . urn , count ) <EOL> count += <NUM_LIT:1> <EOL> def GenerateItems ( self , offset = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> return itertools . islice ( self . _GenerateItems ( ) , offset , self . size ) <EOL> def GetItem ( self , offset = <NUM_LIT:0> ) : <EOL> for item in self . GenerateItems ( offset = offset ) : <EOL> return item <EOL> def __getitem__ ( self , index ) : <EOL> if index >= <NUM_LIT:0> : <EOL> for item in self . GenerateItems ( offset = index ) : <EOL> return item <EOL> else : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> class AFF4Collection ( aff4 . AFF4Volume , RDFValueCollection ) : <EOL> """<STR_LIT>""" <EOL> _rdf_type = rdf_client . AFF4ObjectSummary <EOL> _behaviours = frozenset ( [ "<STR_LIT>" ] ) <EOL> class SchemaCls ( aff4 . AFF4Volume . SchemaCls , RDFValueCollection . SchemaCls ) : <EOL> VIEW = aff4 . Attribute ( "<STR_LIT>" , AFF4CollectionView , <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" , default = "<STR_LIT>" ) <EOL> def CreateView ( self , attributes ) : <EOL> """<STR_LIT>""" <EOL> self . Set ( self . Schema . VIEW ( attributes ) ) <EOL> def Query ( self , filter_string = "<STR_LIT>" , subjects = None , limit = <NUM_LIT:100> ) : <EOL> """<STR_LIT>""" <EOL> if subjects is None : <EOL> subjects = set ( ) <EOL> for obj in self : <EOL> if len ( subjects ) < limit : <EOL> subjects . add ( obj . urn ) <EOL> else : <EOL> break <EOL> else : <EOL> subjects = set ( subjects [ : limit ] ) <EOL> if filter_string : <EOL> ast = aff4 . AFF4QueryParser ( filter_string ) . Parse ( ) <EOL> filter_obj = ast . Compile ( aff4 . AFF4Filter ) <EOL> for subject in aff4 . FACTORY . MultiOpen ( subjects , token = self . token ) : <EOL> if filter_string and not filter_obj . FilterOne ( subject ) : <EOL> continue <EOL> yield subject <EOL> def ListChildren ( self , ** _ ) : <EOL> for aff4object_summary in self : <EOL> yield aff4object_summary . urn <EOL> class GRRSignedBlobCollection ( RDFValueCollection ) : <EOL> _rdf_type = rdf_crypto . SignedBlob <EOL> class GRRSignedBlob ( aff4 . AFF4Stream ) : <EOL> """<STR_LIT>""" <EOL> collection = None <EOL> @ classmethod <EOL> def NewFromContent ( cls , content , urn , chunk_size = <NUM_LIT> , <EOL> token = None , private_key = None , public_key = None , <EOL> prompt = True ) : <EOL> """<STR_LIT>""" <EOL> with aff4 . FACTORY . Create ( <EOL> urn , cls . __name__ , mode = "<STR_LIT:w>" , token = token ) as fd : <EOL> for start_of_chunk in xrange ( <NUM_LIT:0> , len ( content ) , chunk_size ) : <EOL> chunk = content [ start_of_chunk : start_of_chunk + chunk_size ] <EOL> blob_rdf = rdf_crypto . SignedBlob ( ) <EOL> blob_rdf . Sign ( chunk , private_key , public_key , prompt = prompt ) <EOL> fd . Add ( blob_rdf ) <EOL> return urn <EOL> @ property <EOL> def size ( self ) : <EOL> self . _EnsureInitialized ( ) <EOL> return self . _size <EOL> @ size . setter <EOL> def size ( self , value ) : <EOL> self . _size = value <EOL> @ property <EOL> def chunks ( self ) : <EOL> """<STR_LIT>""" <EOL> self . _EnsureInitialized ( ) <EOL> return len ( self . collection ) <EOL> def _EnsureInitialized ( self ) : <EOL> if self . collection is None : <EOL> if self . mode == "<STR_LIT:r>" : <EOL> self . collection = aff4 . FACTORY . Open ( <EOL> self . urn . Add ( "<STR_LIT>" ) , mode = "<STR_LIT:r>" , token = self . token ) <EOL> self . fd = cStringIO . StringIO ( ) <EOL> for x in self . collection : <EOL> self . fd . write ( x . data ) <EOL> self . _size = self . fd . tell ( ) <EOL> self . fd . seek ( <NUM_LIT:0> ) <EOL> elif self . mode == "<STR_LIT:w>" : <EOL> self . fd = cStringIO . StringIO ( ) <EOL> self . _size = <NUM_LIT:0> <EOL> self . collection = aff4 . FACTORY . Create ( <EOL> self . urn . Add ( "<STR_LIT>" ) , "<STR_LIT>" , mode = "<STR_LIT:w>" , <EOL> token = self . token ) <EOL> else : <EOL> raise RuntimeError ( <EOL> "<STR_LIT>" ) <EOL> def Write ( self , length ) : <EOL> raise IOError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> def Read ( self , length ) : <EOL> if self . mode != "<STR_LIT:r>" : <EOL> raise IOError ( "<STR_LIT>" ) <EOL> self . _EnsureInitialized ( ) <EOL> return self . fd . read ( length ) <EOL> def Add ( self , item ) : <EOL> if "<STR_LIT:r>" in self . mode : <EOL> raise IOError ( "<STR_LIT>" ) <EOL> self . _EnsureInitialized ( ) <EOL> self . collection . Add ( item ) <EOL> def __iter__ ( self ) : <EOL> self . _EnsureInitialized ( ) <EOL> return iter ( self . collection ) <EOL> def Close ( self ) : <EOL> super ( GRRSignedBlob , self ) . Close ( ) <EOL> if self . collection is not None : <EOL> self . collection . Close ( ) <EOL> def __len__ ( self ) : <EOL> return self . size <EOL> def Tell ( self ) : <EOL> self . _EnsureInitialized ( ) <EOL> return self . fd . tell ( ) <EOL> def Seek ( self , offset , whence = <NUM_LIT:0> ) : <EOL> self . _EnsureInitialized ( ) <EOL> self . fd . seek ( offset , whence ) <EOL> class GrepResultsCollection ( RDFValueCollection ) : <EOL> """<STR_LIT>""" <EOL> _rdf_type = rdf_client . BufferReference <EOL> class ClientAnomalyCollection ( RDFValueCollection ) : <EOL> """<STR_LIT>""" <EOL> _rdf_type = rdf_anomaly . Anomaly <EOL> class SeekIndexPair ( rdf_structs . RDFProtoStruct ) : <EOL> """<STR_LIT>""" <EOL> protobuf = jobs_pb2 . SeekIndexPair <EOL> class SeekIndex ( rdf_structs . RDFProtoStruct ) : <EOL> """<STR_LIT>""" <EOL> protobuf = jobs_pb2 . SeekIndex <EOL> class PackedVersionedCollection ( RDFValueCollection ) : <EOL> """<STR_LIT>""" <EOL> notification_queue = "<STR_LIT>" <EOL> index_prefix = "<STR_LIT>" <EOL> index_format = index_prefix + "<STR_LIT:%s>" <EOL> @ classmethod <EOL> def ScheduleNotification ( cls , urn , sync = False , token = None ) : <EOL> """<STR_LIT>""" <EOL> data_store . DB . Set ( cls . notification_queue , cls . index_format % urn , <EOL> urn , replace = True , token = token , sync = sync ) <EOL> @ classmethod <EOL> def QueryNotifications ( cls , timestamp = None , token = None ) : <EOL> """<STR_LIT>""" <EOL> if token is None : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if timestamp is None : <EOL> timestamp = rdfvalue . RDFDatetime ( ) . Now ( ) <EOL> for _ , urn , urn_timestamp in data_store . DB . ResolvePrefix ( <EOL> cls . notification_queue , cls . index_prefix , <EOL> timestamp = ( <NUM_LIT:0> , timestamp ) , token = token ) : <EOL> yield rdfvalue . RDFURN ( urn , age = urn_timestamp ) <EOL> @ classmethod <EOL> def DeleteNotifications ( cls , urns , end = None , token = None ) : <EOL> """<STR_LIT>""" <EOL> if token is None : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> predicates = [ cls . index_format % urn for urn in urns ] <EOL> data_store . DB . DeleteAttributes ( cls . notification_queue , predicates , <EOL> end = end , token = token , sync = True ) <EOL> @ classmethod <EOL> def AddToCollection ( cls , collection_urn , rdf_values , sync = True , <EOL> token = None ) : <EOL> """<STR_LIT>""" <EOL> if token is None : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> data_attrs = [ ] <EOL> for rdf_value in rdf_values : <EOL> if rdf_value is None : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if cls . _rdf_type and not isinstance ( rdf_value , cls . _rdf_type ) : <EOL> raise ValueError ( "<STR_LIT>" % <EOL> cls . _rdf_type . __name__ ) <EOL> if not rdf_value . age : <EOL> rdf_value . age . Now ( ) <EOL> data_attrs . append ( cls . SchemaCls . DATA ( <EOL> rdf_protodict . EmbeddedRDFValue ( payload = rdf_value ) ) ) <EOL> attrs_to_set = { cls . SchemaCls . DATA : data_attrs } <EOL> if cls . IsJournalingEnabled ( ) : <EOL> journal_entry = cls . SchemaCls . ADDITION_JOURNAL ( len ( rdf_values ) ) <EOL> attrs_to_set [ cls . SchemaCls . ADDITION_JOURNAL ] = [ journal_entry ] <EOL> aff4 . FACTORY . SetAttributes ( collection_urn , attrs_to_set , set ( ) , <EOL> add_child_index = False , sync = sync , <EOL> token = token ) <EOL> cls . ScheduleNotification ( collection_urn , token = token ) <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" , <EOL> delta = len ( rdf_values ) ) <EOL> class SchemaCls ( RDFValueCollection . SchemaCls ) : <EOL> """<STR_LIT>""" <EOL> DATA = aff4 . Attribute ( "<STR_LIT>" , rdf_protodict . EmbeddedRDFValue , <EOL> "<STR_LIT>" , versioned = True ) <EOL> SEEK_INDEX = aff4 . Attribute ( "<STR_LIT>" , SeekIndex , <EOL> "<STR_LIT>" , versioned = False ) <EOL> ADDITION_JOURNAL = aff4 . Attribute ( "<STR_LIT>" , <EOL> rdfvalue . RDFInteger , <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" , versioned = True ) <EOL> COMPACTION_JOURNAL = aff4 . Attribute ( "<STR_LIT>" , <EOL> rdfvalue . RDFInteger , <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> INDEX_INTERVAL = <NUM_LIT> <EOL> COMPACTION_BATCH_SIZE = <NUM_LIT> <EOL> MAX_REVERSED_RESULTS = <NUM_LIT> <EOL> @ staticmethod <EOL> def IsJournalingEnabled ( ) : <EOL> return config_lib . CONFIG [ <EOL> "<STR_LIT>" ] <EOL> def Flush ( self , sync = True ) : <EOL> send_notification = self . _dirty and self . Schema . DATA in self . new_attributes <EOL> super ( PackedVersionedCollection , self ) . Flush ( sync = sync ) <EOL> if send_notification : <EOL> self . ScheduleNotification ( self . urn , token = self . token ) <EOL> def Close ( self , sync = True ) : <EOL> send_notification = self . _dirty and self . Schema . DATA in self . new_attributes <EOL> super ( PackedVersionedCollection , self ) . Close ( sync = sync ) <EOL> if send_notification : <EOL> self . ScheduleNotification ( self . urn , token = self . token ) <EOL> def Add ( self , rdf_value = None , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if rdf_value is None and self . _rdf_type : <EOL> rdf_value = self . _rdf_type ( ** kwargs ) <EOL> if not rdf_value . age : <EOL> rdf_value . age . Now ( ) <EOL> self . Set ( self . Schema . DATA ( payload = rdf_value ) ) <EOL> if self . IsJournalingEnabled ( ) : <EOL> self . Set ( self . Schema . ADDITION_JOURNAL ( <NUM_LIT:1> ) ) <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" ) <EOL> def AddAll ( self , rdf_values , callback = None ) : <EOL> """<STR_LIT>""" <EOL> for rdf_value in rdf_values : <EOL> if rdf_value is None : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if self . _rdf_type and not isinstance ( rdf_value , self . _rdf_type ) : <EOL> raise ValueError ( "<STR_LIT>" % <EOL> self . _rdf_type . __name__ ) <EOL> if not rdf_value . age : <EOL> rdf_value . age . Now ( ) <EOL> for index , rdf_value in enumerate ( rdf_values ) : <EOL> self . Set ( self . Schema . DATA ( payload = rdf_value ) ) <EOL> if callback : <EOL> callback ( index , rdf_value ) <EOL> if self . IsJournalingEnabled ( ) : <EOL> self . Set ( self . Schema . ADDITION_JOURNAL ( len ( rdf_values ) ) ) <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" , <EOL> delta = len ( rdf_values ) ) <EOL> def GenerateUncompactedItems ( self , max_reversed_results = <NUM_LIT:0> , <EOL> timestamp = None ) : <EOL> if self . IsAttributeSet ( self . Schema . DATA ) : <EOL> freeze_timestamp = timestamp or rdfvalue . RDFDatetime ( ) . Now ( ) <EOL> results = [ ] <EOL> for _ , value , _ in data_store . DB . ResolvePrefix ( <EOL> self . urn , self . Schema . DATA . predicate , token = self . token , <EOL> timestamp = ( <NUM_LIT:0> , freeze_timestamp ) ) : <EOL> if results is not None : <EOL> results . append ( self . Schema . DATA ( value ) . payload ) <EOL> if max_reversed_results and len ( results ) > max_reversed_results : <EOL> for result in results : <EOL> yield result <EOL> results = None <EOL> else : <EOL> yield self . Schema . DATA ( value ) . payload <EOL> if results is not None : <EOL> for result in reversed ( results ) : <EOL> yield result <EOL> def GenerateItems ( self , offset = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> freeze_timestamp = rdfvalue . RDFDatetime ( ) . Now ( ) <EOL> index = <NUM_LIT:0> <EOL> byte_offset = <NUM_LIT:0> <EOL> if offset >= self . INDEX_INTERVAL and self . IsAttributeSet ( <EOL> self . Schema . SEEK_INDEX ) : <EOL> seek_index = self . Get ( self . Schema . SEEK_INDEX ) <EOL> for value in reversed ( seek_index . checkpoints ) : <EOL> if value . index_offset <= offset : <EOL> index = value . index_offset <EOL> byte_offset = value . byte_offset <EOL> break <EOL> for x in self . _GenerateItems ( byte_offset = byte_offset ) : <EOL> if index >= offset : <EOL> yield x <EOL> index += <NUM_LIT:1> <EOL> for x in self . GenerateUncompactedItems ( <EOL> max_reversed_results = self . MAX_REVERSED_RESULTS , <EOL> timestamp = freeze_timestamp ) : <EOL> if index >= offset : <EOL> yield x <EOL> index += <NUM_LIT:1> <EOL> def GetIndex ( self ) : <EOL> """<STR_LIT>""" <EOL> if not self . IsAttributeSet ( self . Schema . SEEK_INDEX ) : <EOL> return [ ] <EOL> else : <EOL> seek_index = self . Get ( self . Schema . SEEK_INDEX ) <EOL> return [ ( v . index_offset , v . byte_offset ) <EOL> for v in reversed ( seek_index . checkpoints ) ] <EOL> @ utils . Synchronized <EOL> def Compact ( self , callback = None , timestamp = None ) : <EOL> """<STR_LIT>""" <EOL> if not self . locked : <EOL> raise aff4 . LockError ( "<STR_LIT>" ) <EOL> compacted_count = <NUM_LIT:0> <EOL> batches_urns = [ ] <EOL> current_batch = [ ] <EOL> freeze_timestamp = timestamp or rdfvalue . RDFDatetime ( ) . Now ( ) <EOL> def UpdateIndex ( ) : <EOL> seek_index = self . Get ( self . Schema . SEEK_INDEX , SeekIndex ( ) ) <EOL> prev_index_pair = seek_index . checkpoints and seek_index . checkpoints [ - <NUM_LIT:1> ] <EOL> if ( not prev_index_pair or <EOL> self . size - prev_index_pair . index_offset >= self . INDEX_INTERVAL ) : <EOL> new_index_pair = SeekIndexPair ( index_offset = self . size , <EOL> byte_offset = self . fd . Tell ( ) ) <EOL> seek_index . checkpoints . Append ( new_index_pair ) <EOL> self . Set ( self . Schema . SEEK_INDEX , seek_index ) <EOL> def DeleteVersionedDataAndFlush ( ) : <EOL> """<STR_LIT>""" <EOL> data_store . DB . DeleteAttributes ( self . urn , [ self . Schema . DATA . predicate ] , <EOL> end = freeze_timestamp , <EOL> token = self . token , sync = True ) <EOL> if self . IsJournalingEnabled ( ) : <EOL> journal_entry = self . Schema . COMPACTION_JOURNAL ( compacted_count , <EOL> age = freeze_timestamp ) <EOL> attrs_to_set = { self . Schema . COMPACTION_JOURNAL : [ journal_entry ] } <EOL> aff4 . FACTORY . SetAttributes ( self . urn , attrs_to_set , set ( ) , <EOL> add_child_index = False , sync = True , <EOL> token = self . token ) <EOL> if self . Schema . DATA in self . synced_attributes : <EOL> del self . synced_attributes [ self . Schema . DATA ] <EOL> self . Flush ( sync = True ) <EOL> def HeartBeat ( ) : <EOL> """<STR_LIT>""" <EOL> lease_time = config_lib . CONFIG [ "<STR_LIT>" ] <EOL> if self . CheckLease ( ) < lease_time / <NUM_LIT:2> : <EOL> logging . info ( "<STR_LIT>" , self . urn ) <EOL> self . UpdateLease ( lease_time ) <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" ) <EOL> if callback : <EOL> callback ( ) <EOL> HeartBeat ( ) <EOL> for _ , value , _ in data_store . DB . ResolvePrefix ( <EOL> self . urn , self . Schema . DATA . predicate , token = self . token , <EOL> timestamp = ( <NUM_LIT:0> , freeze_timestamp ) ) : <EOL> HeartBeat ( ) <EOL> current_batch . append ( value ) <EOL> compacted_count += <NUM_LIT:1> <EOL> if len ( current_batch ) >= self . COMPACTION_BATCH_SIZE : <EOL> batch_urn = rdfvalue . RDFURN ( "<STR_LIT>" ) . Add ( <EOL> "<STR_LIT>" % utils . PRNG . GetULong ( ) ) <EOL> batches_urns . append ( batch_urn ) <EOL> buf = cStringIO . StringIO ( ) <EOL> for data in reversed ( current_batch ) : <EOL> buf . write ( struct . pack ( "<STR_LIT>" , len ( data ) ) ) <EOL> buf . write ( data ) <EOL> with aff4 . FACTORY . Create ( batch_urn , "<STR_LIT>" , mode = "<STR_LIT:w>" , <EOL> token = self . token ) as batch_stream : <EOL> batch_stream . Write ( buf . getvalue ( ) ) <EOL> current_batch = [ ] <EOL> if not current_batch and not batches_urns : <EOL> return <NUM_LIT:0> <EOL> if current_batch : <EOL> buf = cStringIO . StringIO ( ) <EOL> for data in reversed ( current_batch ) : <EOL> buf . write ( struct . pack ( "<STR_LIT>" , len ( data ) ) ) <EOL> buf . write ( data ) <EOL> self . fd . Seek ( <NUM_LIT:0> , <NUM_LIT:2> ) <EOL> self . fd . Write ( buf . getvalue ( ) ) <EOL> self . stream_dirty = True <EOL> self . size += len ( current_batch ) <EOL> UpdateIndex ( ) <EOL> if not batches_urns : <EOL> DeleteVersionedDataAndFlush ( ) <EOL> return compacted_count <EOL> batches = { } <EOL> for batch in aff4 . FACTORY . MultiOpen ( batches_urns , aff4_type = "<STR_LIT>" , <EOL> token = self . token ) : <EOL> batches [ batch . urn ] = batch <EOL> if len ( batches_urns ) != len ( batches ) : <EOL> raise RuntimeError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> self . fd . Seek ( <NUM_LIT:0> , <NUM_LIT:2> ) <EOL> for batch_urn in reversed ( batches_urns ) : <EOL> batch = batches [ batch_urn ] <EOL> HeartBeat ( ) <EOL> data = batch . Read ( len ( batch ) ) <EOL> self . fd . Write ( data ) <EOL> self . stream_dirty = True <EOL> self . size += self . COMPACTION_BATCH_SIZE <EOL> UpdateIndex ( ) <EOL> aff4 . FACTORY . Delete ( batch_urn , token = self . token ) <EOL> DeleteVersionedDataAndFlush ( ) <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" , <EOL> delta = compacted_count ) <EOL> return compacted_count <EOL> def CalculateLength ( self ) : <EOL> length = super ( PackedVersionedCollection , self ) . __len__ ( ) <EOL> if self . IsAttributeSet ( self . Schema . DATA ) : <EOL> if self . age_policy == aff4 . ALL_TIMES : <EOL> length += len ( list ( self . GetValuesForAttribute ( self . Schema . DATA ) ) ) <EOL> else : <EOL> length += len ( list ( data_store . DB . ResolveMulti ( <EOL> self . urn , [ self . Schema . DATA . predicate ] , token = self . token , <EOL> timestamp = data_store . DB . ALL_TIMESTAMPS ) ) ) <EOL> return length <EOL> def __len__ ( self ) : <EOL> return self . CalculateLength ( ) <EOL> def __nonzero__ ( self ) : <EOL> if "<STR_LIT:r>" not in self . mode : <EOL> raise AttributeError ( <EOL> "<STR_LIT>" ) <EOL> if super ( PackedVersionedCollection , self ) . __nonzero__ ( ) : <EOL> return True <EOL> return self . IsAttributeSet ( self . Schema . DATA ) <EOL> class ResultsOutputCollection ( PackedVersionedCollection ) : <EOL> """<STR_LIT>""" <EOL> notification_queue = "<STR_LIT>" <EOL> class SchemaCls ( PackedVersionedCollection . SchemaCls ) : <EOL> RESULTS_SOURCE = aff4 . Attribute ( "<STR_LIT>" , rdfvalue . RDFURN , <EOL> "<STR_LIT>" ) <EOL> def Initialize ( self ) : <EOL> super ( ResultsOutputCollection , self ) . Initialize ( ) <EOL> if "<STR_LIT:w>" in self . mode and self . fd . size == <NUM_LIT:0> : <EOL> self . fd . SetChunksize ( <NUM_LIT> * <NUM_LIT> ) <EOL> class CollectionsInitHook ( registry . InitHook ) : <EOL> def RunOnce ( self ) : <EOL> """<STR_LIT>""" <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) </s>
