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)
