def testCorrectlyExportsFourValuesOfTwoDifferentTypes(self): with self.pool: self.collection.Add(rdf_flows.GrrMessage( payload=rdfvalue.RDFString("foo"), source=self.client_id), mutation_pool=self.pool) self.collection.Add(rdf_flows.GrrMessage( payload=rdfvalue.RDFInteger(42), source=self.client_id), mutation_pool=self.pool) self.collection.Add(rdf_flows.GrrMessage( payload=rdfvalue.RDFString("bar"), source=self.client_id), mutation_pool=self.pool) self.collection.Add(rdf_flows.GrrMessage( payload=rdfvalue.RDFInteger(43), source=self.client_id), mutation_pool=self.pool) chunks = self.ProcessPlugin() self.assertListEqual(chunks, [ "Start: aff4:/foo/bar", "Values of type: RDFInteger", "First pass: 42 (source=%s)" % self.client_id, "First pass: 43 (source=%s)" % self.client_id, "Second pass: 42 (source=%s)" % self.client_id, "Second pass: 43 (source=%s)" % self.client_id, "Values of type: RDFString", "First pass: foo (source=%s)" % self.client_id, "First pass: bar (source=%s)" % self.client_id, "Second pass: foo (source=%s)" % self.client_id, "Second pass: bar (source=%s)" % self.client_id, "Finish: aff4:/foo/bar" ]) # pyformat: disable
def TestAnd(val1, val2, expected): self.assertEqual(rdfvalue.RDFInteger(val1) & val2, expected) self.assertEqual(val1 & rdfvalue.RDFInteger(val2), expected) value = rdfvalue.RDFInteger(val1) value &= val2 self.assertEqual(value, expected) value = val1 value &= rdfvalue.RDFInteger(val2) self.assertEqual(value, expected)
def TestOr(val1, val2, expected): self.assertEqual(rdfvalue.RDFInteger(val1) | val2, expected) self.assertEqual(val1 | rdfvalue.RDFInteger(val2), expected) value = rdfvalue.RDFInteger(val1) value |= val2 self.assertEqual(value, expected) value = val1 value |= rdfvalue.RDFInteger(val2) self.assertEqual(value, expected)
def testAutoIndexing(self): indexing_done = threading.Event() def UpdateIndex(_): indexing_done.set() # To reduce the time for the test to run, reduce the delays, so that # indexing should happen instantaneously. isq = sequential_collection.IndexedSequentialCollection biu = sequential_collection.BACKGROUND_INDEX_UPDATER with utils.MultiStubber( (biu, "INDEX_DELAY", 0), (isq, "INDEX_WRITE_DELAY", rdfvalue.Duration("0s")), (isq, "INDEX_SPACING", 8), (isq, "UpdateIndex", UpdateIndex)): urn = "aff4:/sequential_collection/testAutoIndexing" collection = self._TestCollection(urn) # TODO(amoser): Without using a mutation pool, this test is really # slow on MySQL data store. with data_store.DB.GetMutationPool() as pool: for i in range(2048): collection.StaticAdd(rdfvalue.RDFURN(urn), rdfvalue.RDFInteger(i), mutation_pool=pool) # Wait for the updater thread to finish the indexing. if not indexing_done.wait(timeout=10): self.fail("Indexing did not finish in time.")
def testMessageHandlerRequests(self): requests = [ rdf_objects.MessageHandlerRequest( client_id="C.1000000000000000", handler_name="Testhandler", request_id=i * 100, request=rdfvalue.RDFInteger(i)) for i in range(5) ] self.db.WriteMessageHandlerRequests(requests) read = self.db.ReadMessageHandlerRequests() for r in read: self.assertTrue(r.timestamp) r.timestamp = None self.assertEqual(sorted(read, key=lambda req: req.request_id), requests) self.db.DeleteMessageHandlerRequests(requests[:2]) self.db.DeleteMessageHandlerRequests(requests[4]) read = self.db.ReadMessageHandlerRequests() self.assertEqual(len(read), 2) read = sorted(read, key=lambda req: req.request_id) for r in read: r.timestamp = None self.assertEqual(requests[2:4], read) self.db.DeleteMessageHandlerRequests(read)
def testClaimReturnsPreviouslyClaimedRecordsAfterTimeout(self): queue_urn = ( "aff4:/queue_test/testClaimReturnsPreviouslyClaimedRecordsAfterTimeout" ) with self.pool: with aff4.FACTORY.Create(queue_urn, TestQueue, token=self.token) as queue: for i in range(100): queue.Add(rdfvalue.RDFInteger(i), mutation_pool=self.pool) data_store.DB.Flush() with aff4.FACTORY.OpenWithLock(queue_urn, lease_time=200, token=self.token) as queue: results_1 = queue.ClaimRecords() self.assertEqual(100, len(results_1)) with test_lib.FakeTime(rdfvalue.RDFDatetime.Now() + rdfvalue.Duration("45m")): with aff4.FACTORY.OpenWithLock(queue_urn, lease_time=200, token=self.token) as queue: results_2 = queue.ClaimRecords() self.assertEqual(100, len(results_2))
def testClaimReturnsPreviouslyReleasedRecords(self): queue_urn = "aff4:/queue_test/testClaimReturnsPreviouslyReleasedRecords" with self.pool: with aff4.FACTORY.Create(queue_urn, TestQueue, token=self.token) as queue: for i in range(100): queue.Add(rdfvalue.RDFInteger(i), mutation_pool=self.pool) data_store.DB.Flush() with aff4.FACTORY.OpenWithLock(queue_urn, lease_time=200, token=self.token) as queue: results = queue.ClaimRecords() odd_ids = [ record for record in results if int(record.value) % 2 == 1 ] queue.ReleaseRecord(odd_ids[0], token=self.token) queue.ReleaseRecords(odd_ids[1:], token=self.token) with aff4.FACTORY.OpenWithLock(queue_urn, lease_time=200, token=self.token) as queue: odd_results = queue.ClaimRecords() self.assertEqual(50, len(odd_results)) self.assertEqual(1, odd_results[0].value) self.assertEqual(99, odd_results[49].value)
def End(self, responses): self.SendReply(rdfvalue.RDFInteger(42)) self.SendReply(rdfvalue.RDFString("foo bar")) self.SendReply(rdfvalue.RDFString("foo1 bar1")) self.SendReply(rdfvalue.RDFURN("aff4:/foo/bar")) self.SendReply(rdfvalue.RDFURN("aff4:/foo1/bar1")) self.SendReply(rdfvalue.RDFURN("aff4:/foo2/bar2"))
def testWellKnownFlows(self): """Test the well known flows.""" test_flow = self.FlowSetup(flow_test_lib.WellKnownSessionTest.__name__) # Make sure the session ID is well known self.assertEqual( test_flow.session_id, flow_test_lib.WellKnownSessionTest.well_known_session_id) # Messages to Well Known flows can be unauthenticated messages = [ rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(i)) for i in range(10) ] for message in messages: test_flow.ProcessMessage(message) # The messages might be processed in arbitrary order test_flow.messages.sort() # Make sure that messages were processed even without a status # message to complete the transaction (Well known flows do not # have transactions or states - all messages always get to the # ProcessMessage method): self.assertEqual(test_flow.messages, range(10))
def testFlowStoresResultsPerType(self): flow_urn = flow_test_lib.TestFlowHelper( FlowWithMultipleResultTypes.__name__, action_mocks.ActionMock(), token=self.token, client_id=self.client_id) c = flow.GRRFlow.TypedResultCollectionForFID(flow_urn) self.assertEqual( set(c.ListStoredTypes()), set([ rdfvalue.RDFInteger.__name__, rdfvalue.RDFString.__name__, rdfvalue.RDFURN.__name__ ])) self.assertEqual(c.LengthByType(rdfvalue.RDFInteger.__name__), 1) self.assertEqual(c.LengthByType(rdfvalue.RDFString.__name__), 2) self.assertEqual(c.LengthByType(rdfvalue.RDFURN.__name__), 3) self.assertListEqual( [v.payload for _, v in c.ScanByType(rdfvalue.RDFInteger.__name__)], [rdfvalue.RDFInteger(42)]) self.assertListEqual( [v.payload for _, v in c.ScanByType(rdfvalue.RDFString.__name__)], [rdfvalue.RDFString("foo bar"), rdfvalue.RDFString("foo1 bar1")]) self.assertListEqual( [v.payload for _, v in c.ScanByType(rdfvalue.RDFURN.__name__)], [ rdfvalue.RDFURN("foo/bar"), rdfvalue.RDFURN("foo1/bar1"), rdfvalue.RDFURN("foo2/bar2") ])
def testDeleteRemovesRecords(self): queue_urn = "aff4:/queue_test/testDeleteRemovesRecords" with self.pool: with aff4.FACTORY.Create(queue_urn, TestQueue, token=self.token) as queue: for i in range(100): queue.Add(rdfvalue.RDFInteger(i), mutation_pool=self.pool) data_store.DB.Flush() with aff4.FACTORY.OpenWithLock(queue_urn, lease_time=200, token=self.token) as queue: results = queue.ClaimRecords() with aff4.FACTORY.OpenWithLock(queue_urn, lease_time=200, token=self.token) as queue: queue.DeleteRecord(results[0], token=self.token) queue.DeleteRecords(results[1:], token=self.token) # Wait past the default claim length, to make sure that delete actually did # something. with test_lib.FakeTime(rdfvalue.RDFDatetime.Now() + rdfvalue.Duration("45m")): with aff4.FACTORY.OpenWithLock(queue_urn, lease_time=200, token=self.token) as queue: results = queue.ClaimRecords() self.assertEqual(0, len(results))
def testWrapsValueInGrrMessageIfNeeded(self): with self.pool: self.collection.Add(rdfvalue.RDFInteger(42), mutation_pool=self.pool) items = list(self.collection) self.assertTrue(isinstance(items[0], rdf_flows.GrrMessage)) self.assertEqual(items[0].payload, 42)
def testIndexCreate(self): spacing = 10 with utils.Stubber(sequential_collection.IndexedSequentialCollection, "INDEX_SPACING", spacing): urn = "aff4:/sequential_collection/testIndexCreate" collection = self._TestCollection(urn) # TODO(amoser): Without using a mutation pool, this test is really # slow on MySQL data store. with data_store.DB.GetMutationPool() as pool: for i in range(10 * spacing): collection.StaticAdd(urn, rdfvalue.RDFInteger(i), mutation_pool=pool) # It is too soon to build an index, check that we don't. self.assertEqual(collection._index, None) self.assertEqual(collection.CalculateLength(), 10 * spacing) self.assertEqual(sorted(collection._index.keys()), [0]) now = time.time() * 1e6 twenty_seconds_ago = (time.time() - 20) * 1e6 # Push the clock forward 10m, and we should build an index on access. with test_lib.FakeTime(rdfvalue.RDFDatetime.Now() + rdfvalue.Duration("10m")): # Read from start doesn't rebuild index (lazy rebuild) _ = collection[0] self.assertEqual(sorted(collection._index.keys()), [0]) self.assertEqual(collection.CalculateLength(), 10 * spacing) self.assertEqual(sorted(collection._index.keys()), [i * spacing for i in xrange(10)]) for index in collection._index: if not index: continue timestamp, suffix = collection._index[index] self.assertLessEqual(twenty_seconds_ago, timestamp) self.assertLessEqual(timestamp, now) self.assertTrue(0 <= suffix <= 0xFFFFFF) # Now check that the index was persisted to aff4 by re-opening # and checking that a read from head does load full index # (optimistic load): collection = self._TestCollection( "aff4:/sequential_collection/testIndexCreate") self.assertEqual(collection._index, None) _ = collection[0] self.assertEqual(sorted(collection._index.keys()), [i * spacing for i in xrange(10)]) for index in collection._index: if not index: continue timestamp, suffix = collection._index[index] self.assertLessEqual(twenty_seconds_ago, timestamp) self.assertLessEqual(timestamp, now) self.assertTrue(0 <= suffix <= 0xFFFFFF)
def testValuesOfSingleTypeAreAddedAndIterated(self): with self.pool: for i in range(100): self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(i)), mutation_pool=self.pool) for index, v in enumerate(self.collection): self.assertEqual(index, v.payload)
def testValuesOfMultipleTypesCanBeIteratedTogether(self): with self.pool: original_values = set() for i in range(100): self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(i)), mutation_pool=self.pool) original_values.add( rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(i))) self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFString(i)), mutation_pool=self.pool) original_values.add( rdf_flows.GrrMessage(payload=rdfvalue.RDFString(i))) self.assertEqual(sorted([v.payload for v in original_values]), sorted([v.payload for v in self.collection]))
def testLengthOfCollectionIsCorrectWhenMultipleTypesAreUsed(self): with self.pool: for i in range(100): self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(i)), mutation_pool=self.pool) self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFString(i)), mutation_pool=self.pool) self.assertEqual(200, len(self.collection))
def testAddGet(self): collection = sequential_collection.GeneralIndexedCollection( rdfvalue.RDFURN("aff4:/sequential_collection/testAddGetIndexed")) with data_store.DB.GetMutationPool() as pool: collection.Add(rdfvalue.RDFInteger(42), mutation_pool=pool) collection.Add(rdfvalue.RDFString("the meaning of life"), mutation_pool=pool) self.assertEqual(collection[0].__class__, rdfvalue.RDFInteger) self.assertEqual(collection[0], 42) self.assertEqual(collection[1].__class__, rdfvalue.RDFString) self.assertEqual(collection[1], "the meaning of life")
def testAddGet(self): collection = self._TestCollection( "aff4:/sequential_collection/testAddGet") self.assertEqual(collection.CalculateLength(), 0) with data_store.DB.GetMutationPool() as pool: for i in range(100): collection.Add(rdfvalue.RDFInteger(i), mutation_pool=pool) for i in range(100): self.assertEqual(collection[i], i) self.assertEqual(collection.CalculateLength(), 100) self.assertEqual(len(collection), 100)
def testSimpleTypeAssignment(self): sample = rdf_client.StatEntry() sample.AddDescriptor( rdf_structs.ProtoRDFValue(name="test", field_number=45, default=rdfvalue.RDFInteger(0), rdf_type=rdfvalue.RDFInteger)) self.assertIsInstance(sample.test, rdfvalue.RDFInteger) # Can we assign an RDFValue instance? sample.test = rdfvalue.RDFInteger(5) self.assertEqual(sample.test, 5) # Check that bare values can be coerced. sample.test = 6 self.assertIsInstance(sample.test, rdfvalue.RDFInteger) self.assertEqual(sample.test, 6) # Assign an enum. sample.registry_type = sample.RegistryType.REG_DWORD self.assertEqual(sample.registry_type, sample.RegistryType.REG_DWORD) sample.registry_type = rdf_client.StatEntry.RegistryType.REG_SZ self.assertEqual(sample.registry_type, sample.RegistryType.REG_SZ) # We can also assign the string value. sample.registry_type = "REG_QWORD" self.assertEqual(sample.registry_type, sample.RegistryType.REG_QWORD) # Check that coercing works. sample.test = "10" self.assertEqual(sample.test, 10) # Assign an RDFValue which can not be coerced. self.assertRaises(type_info.TypeValueError, setattr, sample, "test", rdfvalue.RDFString("hello"))
def testDelete(self): collection = self._TestCollection( "aff4:/sequential_collection/testDelete") with data_store.DB.GetMutationPool() as pool: for i in range(100): collection.Add(rdfvalue.RDFInteger(i), mutation_pool=pool) collection.Delete() collection = self._TestCollection( "aff4:/sequential_collection/testDelete") for _ in collection.Scan(): self.fail( "Deleted and recreated SequentialCollection should be empty")
def testAddScan(self): collection = self._TestCollection( "aff4:/sequential_collection/testAddScan") with data_store.DB.GetMutationPool() as pool: for i in range(100): collection.Add(rdfvalue.RDFInteger(i), mutation_pool=pool) i = 0 last_ts = 0 for (ts, v) in collection.Scan(): last_ts = ts self.assertEqual(i, v) i += 1 with data_store.DB.GetMutationPool() as pool: for j in range(100): collection.Add(rdfvalue.RDFInteger(j + 100), mutation_pool=pool) for (ts, v) in collection.Scan(after_timestamp=last_ts): self.assertEqual(i, v) i += 1 self.assertEqual(i, 200)
def testListing(self): test_urn = "aff4:/sequential_collection/testIndexedListing" collection = self._TestCollection(test_urn) timestamps = [] with data_store.DB.GetMutationPool() as pool: for i in range(100): timestamps.append( collection.Add(rdfvalue.RDFInteger(i), mutation_pool=pool)) with test_lib.Instrument(sequential_collection.SequentialCollection, "Scan") as scan: self.assertEqual(len(list(collection)), 100) # Listing should be done using a single scan but there is another one # for calculating the length. self.assertEqual(scan.call_count, 2)
def testDuplicateTimestamps(self): collection = self._TestCollection( "aff4:/sequential_collection/testDuplicateTimestamps") t = rdfvalue.RDFDatetime.Now() with data_store.DB.GetMutationPool() as pool: for i in range(10): ts = collection.Add(rdfvalue.RDFInteger(i), timestamp=t, mutation_pool=pool) self.assertEqual(ts[0], t) i = 0 for (ts, _) in collection.Scan(): self.assertEqual(ts, t) i += 1 self.assertEqual(i, 10)
def testLengthIsReportedCorrectlyForEveryType(self): with self.pool: for i in range(99): self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(i)), mutation_pool=self.pool) for i in range(101): self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFString(i)), mutation_pool=self.pool) self.assertEqual( 99, self.collection.LengthByType(rdfvalue.RDFInteger.__name__)) self.assertEqual( 101, self.collection.LengthByType(rdfvalue.RDFString.__name__))
def testValuesOfMultipleTypesCanBeIteratedPerType(self): with self.pool: for i in range(100): self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(i)), mutation_pool=self.pool) self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFString(i)), mutation_pool=self.pool) for index, (_, v) in enumerate( self.collection.ScanByType(rdfvalue.RDFInteger.__name__)): self.assertEqual(index, v.payload) for index, (_, v) in enumerate( self.collection.ScanByType(rdfvalue.RDFString.__name__)): self.assertEqual(str(index), v.payload)
def testExtractsTypesFromGrrMessage(self): with self.pool: self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(0)), mutation_pool=self.pool) self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFString("foo")), mutation_pool=self.pool) self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFURN("aff4:/foo/bar")), mutation_pool=self.pool) self.assertEqual( set([ rdfvalue.RDFInteger.__name__, rdfvalue.RDFString.__name__, rdfvalue.RDFURN.__name__ ]), set(self.collection.ListStoredTypes()))
def testClaimReturnsRecordsInOrder(self): queue_urn = "aff4:/queue_test/testClaimReturnsRecordsInOrder" with self.pool: with aff4.FACTORY.Create(queue_urn, TestQueue, token=self.token) as queue: for i in range(100): queue.Add(rdfvalue.RDFInteger(i), mutation_pool=self.pool) data_store.DB.Flush() with aff4.FACTORY.OpenWithLock(queue_urn, lease_time=200, token=self.token) as queue: results = queue.ClaimRecords() self.assertEqual(100, len(results)) self.assertEqual(0, results[0].value) self.assertEqual(99, results[99].value)
def testClaimFiltersByStartTime(self): queue_urn = "aff4:/queue_test/testClaimFiltersByStartTime" middle = None with self.pool: with aff4.FACTORY.Create(queue_urn, TestQueue, token=self.token) as queue: for i in range(100): if i == 50: middle = rdfvalue.RDFDatetime.Now() queue.Add(rdfvalue.RDFInteger(i), mutation_pool=self.pool) with aff4.FACTORY.OpenWithLock(queue_urn, lease_time=200, token=self.token) as queue: results = queue.ClaimRecords(start_time=middle) self.assertEqual(50, len(results)) self.assertEqual(50, results[0].value)
def testMessageHandlerRequestSorting(self): for i, ts in enumerate([ 10000, 11000, 12000, 13000, 14000, 19000, 18000, 17000, 16000, 15000 ]): with test_lib.FakeTime( rdfvalue.RDFDatetime.FromSecondsSinceEpoch(ts)): request = rdf_objects.MessageHandlerRequest( client_id="C.1000000000000000", handler_name="Testhandler", request_id=i * 100, request=rdfvalue.RDFInteger(i)) self.db.WriteMessageHandlerRequests([request]) read = self.db.ReadMessageHandlerRequests() for i in range(1, len(read)): self.assertGreater(read[i - 1].timestamp, read[i].timestamp)
def testDeletingCollectionDeletesAllSubcollections(self): if not isinstance(data_store.DB, fake_data_store.FakeDataStore): self.skipTest("Only supported on FakeDataStore.") with self.pool: self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFInteger(0)), mutation_pool=self.pool) self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFString("foo")), mutation_pool=self.pool) self.collection.Add( rdf_flows.GrrMessage(payload=rdfvalue.RDFURN("aff4:/foo/bar")), mutation_pool=self.pool) self.collection.Delete() for urn in data_store.DB.subjects.keys(): self.assertFalse( utils.SmartStr(self.collection.collection_id) in urn)