def testAgePreserved(self):
data = rdfvalue.RDFValueArray([1, 2, 3])
now = rdfvalue.RDFDatetime().Now()
original_age = data.age.Now()
self.assertTrue((now - data.age) < rdfvalue.Duration("5s"))
embedded = rdfvalue.EmbeddedRDFValue(payload=data)
self.assertEqual(embedded.payload.age, original_age)
new_log = rdfvalue.EmbeddedRDFValue(embedded).payload
self.assertEqual(new_log.age, original_age, "Age not preserved: %s != %s" %
(new_log.age.AsMicroSecondsFromEpoch(),
original_age.AsMicroSecondsFromEpoch()))
def AddAll(self, rdf_values, callback=None):
"""Adds a list of rdfvalues to the collection."""
for rdf_value in rdf_values:
if rdf_value is None:
raise ValueError(
"Can't add None to the collection via AddAll.")
if self._rdf_type and not isinstance(rdf_value, self._rdf_type):
raise ValueError(
"This collection only accepts values of type %s" %
self._rdf_type.__name__)
if not rdf_value.age:
rdf_value.age.Now()
buf = cStringIO.StringIO()
for index, rdf_value in enumerate(rdf_values):
data = rdfvalue.EmbeddedRDFValue(
payload=rdf_value).SerializeToString()
buf.write(struct.pack("<i", len(data)))
buf.write(data)
self.size += 1
if callback:
callback(index, rdf_value)
self.fd.Seek(0, 2)
self.fd.Write(buf.getvalue())
self.stream_dirty = True
def AddToCollection(cls, collection_urn, rdf_values, sync=True,
token=None):
"""Adds RDFValues to the collection with a given urn."""
if token is None:
raise ValueError("Token can't be None.")
data_attrs = []
for rdf_value in rdf_values:
if rdf_value is None:
raise ValueError("Can't add None to the collection.")
if cls._rdf_type and not isinstance(rdf_value, cls._rdf_type):
raise ValueError("This collection only accepts values of type %s" %
cls._rdf_type.__name__)
if not rdf_value.age:
rdf_value.age.Now()
data_attrs.append(cls.SchemaCls.DATA(
rdfvalue.EmbeddedRDFValue(payload=rdf_value)))
attrs_to_set = {cls.SchemaCls.DATA: data_attrs}
if cls.IsJournalingEnabled():
journal_entry = cls.SchemaCls.ADDITION_JOURNAL(len(rdf_values))
attrs_to_set[cls.SchemaCls.ADDITION_JOURNAL] = [journal_entry]
aff4.FACTORY.SetAttributes(collection_urn, attrs_to_set, set(),
add_child_index=False, sync=sync,
token=token)
cls.ScheduleNotification(collection_urn, token=token)
# Update system-wide stats.
stats.STATS.IncrementCounter("packed_collection_added",
delta=len(rdf_values))
def GenerateItems(self, offset=0):
"""Iterate over all contained RDFValues.
Args:
offset: The offset in the stream to start reading from.
Yields:
RDFValues stored in the collection.
Raises:
RuntimeError: if we are in write mode.
"""
if not self.fd:
return
if self.mode == "w":
raise RuntimeError("Can not read when in write mode.")
self.fd.seek(offset)
count = 0
while True:
offset = self.fd.Tell()
try:
length = struct.unpack("<i", self.fd.Read(4))[0]
serialized_event = self.fd.Read(length)
except struct.error:
break
result = rdfvalue.EmbeddedRDFValue(serialized_event)
payload = result.payload
if payload is not None:
# Mark the RDFValue with important information relating to the
# collection it is from.
payload.id = count
payload.collection_offset = offset
yield payload
else:
logging.warning(
"payload=None was encountered in a collection %s "
"(index %d), this may mean a logical bug or corrupt "
"data. Ignoring...", self.urn, count)
count += 1
def Add(self, rdf_value=None, **kwargs):
"""Add the rdf value to the collection."""
if rdf_value is None:
if self._rdf_type:
rdf_value = self._rdf_type(**kwargs) # pylint: disable=not-callable
else:
raise ValueError("RDFValueCollection doesn't accept None values.")
if self._rdf_type and not isinstance(rdf_value, self._rdf_type):
raise ValueError("This collection only accepts values of type %s" %
self._rdf_type.__name__)
if not rdf_value.age:
rdf_value.age.Now()
data = rdfvalue.EmbeddedRDFValue(payload=rdf_value).SerializeToString()
self.fd.Seek(0, 2)
self.fd.Write(struct.pack("<i", len(data)))
self.fd.Write(data)
self.size += 1
self._dirty = True
def testChunkSize(self):
urn = "aff4:/test/chunktest"
fd = aff4.FACTORY.Create(urn,
"RDFValueCollection",
mode="w",
token=self.token)
fd.SetChunksize(1024 * 1024)
# Estimate the size of the resulting message.
msg = rdfvalue.GrrMessage(request_id=100)
msg_size = len(
rdfvalue.EmbeddedRDFValue(payload=msg).SerializeToString())
# Write ~500Kb.
n = 500 * 1024 / msg_size
fd.AddAll([rdfvalue.GrrMessage(request_id=i) for i in xrange(n)])
self.assertEqual(fd.fd.Get(fd.fd.Schema._CHUNKSIZE), 1024 * 1024)
# There should be 500K of data.
self.assertGreater(fd.fd.size, 400 * 1024)
# and there should only be one chunk since 500K is less than the chunk size.
self.assertEqual(len(fd.fd.chunk_cache._hash), 1)
fd.Close()
# Closing the collection empties the chunk_cache.
self.assertEqual(len(fd.fd.chunk_cache._hash), 0)
self.assertRaises(ValueError, fd.SetChunksize, (10))
fd = aff4.FACTORY.Open(urn,
"RDFValueCollection",
mode="rw",
token=self.token)
self.assertRaises(ValueError, fd.SetChunksize, (2 * 1024 * 1024))
def GenerateSample(self, number=0):
return rdfvalue.EmbeddedRDFValue(rdfvalue.RDFValueArray([number]))
