<s> """<STR_LIT>""" <EOL> from grr . lib import aff4 <EOL> from grr . lib import data_store <EOL> from grr . lib import rdfvalue <EOL> class AFF4KeywordIndex ( aff4 . AFF4Object ) : <EOL> """<STR_LIT>""" <EOL> INDEX_PREFIX = "<STR_LIT>" <EOL> INDEX_PREFIX_LEN = len ( INDEX_PREFIX ) <EOL> INDEX_COLUMN_FORMAT = INDEX_PREFIX + "<STR_LIT:%s>" <EOL> FIRST_TIMESTAMP = <NUM_LIT:0> <EOL> LAST_TIMESTAMP = ( <NUM_LIT:2> ** <NUM_LIT> ) - <NUM_LIT:2> <EOL> def _KeywordToURN ( self , keyword ) : <EOL> return self . urn . Add ( keyword ) <EOL> def Lookup ( self , keywords , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> posting_lists = self . ReadPostingLists ( keywords , ** kwargs ) <EOL> results = posting_lists . values ( ) <EOL> relevant_set = results [ <NUM_LIT:0> ] <EOL> for hits in results : <EOL> relevant_set &= hits <EOL> if not relevant_set : <EOL> return relevant_set <EOL> return relevant_set <EOL> def ReadPostingLists ( self , keywords , start_time = FIRST_TIMESTAMP , <EOL> end_time = LAST_TIMESTAMP , last_seen_map = None ) : <EOL> """<STR_LIT>""" <EOL> keyword_urns = { self . _KeywordToURN ( k ) : k for k in keywords } <EOL> result = { } <EOL> for kw in keywords : <EOL> result [ kw ] = set ( ) <EOL> for keyword_urn , value in data_store . DB . MultiResolvePrefix ( <EOL> keyword_urns . keys ( ) , self . INDEX_PREFIX , <EOL> timestamp = ( start_time , end_time + <NUM_LIT:1> ) , token = self . token ) : <EOL> for column , _ , ts in value : <EOL> kw = keyword_urns [ keyword_urn ] <EOL> name = column [ self . INDEX_PREFIX_LEN : ] <EOL> result [ kw ] . add ( name ) <EOL> if last_seen_map is not None : <EOL> last_seen_map [ ( kw , name ) ] = max ( last_seen_map . get ( ( kw , name ) , - <NUM_LIT:1> ) , ts ) <EOL> return result <EOL> def AddKeywordsForName ( self , name , keywords , sync = True , timestamp = None , <EOL> ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if timestamp is None : <EOL> timestamp = rdfvalue . RDFDatetime ( ) . Now ( ) . AsMicroSecondsFromEpoch ( ) <EOL> if sync : <EOL> with data_store . DB . GetMutationPool ( token = self . token ) as mutation_pool : <EOL> for keyword in set ( keywords ) : <EOL> mutation_pool . Set ( <EOL> self . _KeywordToURN ( keyword ) , <EOL> self . INDEX_COLUMN_FORMAT % name , "<STR_LIT>" , <EOL> timestamp = timestamp , ** kwargs ) <EOL> else : <EOL> for keyword in set ( keywords ) : <EOL> data_store . DB . Set ( <EOL> self . _KeywordToURN ( keyword ) , <EOL> self . INDEX_COLUMN_FORMAT % name , "<STR_LIT>" , <EOL> token = self . token , sync = False , timestamp = timestamp , ** kwargs ) <EOL> def RemoveKeywordsForName ( self , name , keywords , sync = True ) : <EOL> """<STR_LIT>""" <EOL> if sync : <EOL> with data_store . DB . GetMutationPool ( token = self . token ) as mutation_pool : <EOL> for keyword in set ( keywords ) : <EOL> mutation_pool . DeleteAttributes ( <EOL> self . _KeywordToURN ( keyword ) , <EOL> [ self . INDEX_COLUMN_FORMAT % name ] ) <EOL> else : <EOL> for keyword in set ( keywords ) : <EOL> data_store . DB . DeleteAttributes ( <EOL> self . _KeywordToURN ( keyword ) , <EOL> [ self . INDEX_COLUMN_FORMAT % name ] , <EOL> token = self . token , sync = False ) </s>
