def testAgePreserved(self):
data = rdf_protodict.RDFValueArray([1, 2, 3])
now = rdfvalue.RDFDatetime().Now()
original_age = data.age.Now()
self.assertTrue((now - data.age) < rdfvalue.Duration("5s"))
embedded = rdf_protodict.EmbeddedRDFValue(payload=data)
self.assertEqual(embedded.payload.age, original_age)
new_log = rdf_protodict.EmbeddedRDFValue(embedded).payload
self.assertEqual(
new_log.age, original_age, "Age not preserved: %s != %s" %
(new_log.age.AsMicroSecondsFromEpoch(),
original_age.AsMicroSecondsFromEpoch()))
def _GenerateItems(self, byte_offset=0):
"""Generates items starting from a given byte offset."""
if not self.fd:
return
if self.mode == "w":
raise RuntimeError("Can not read when in write mode.")
self.fd.seek(byte_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 = rdf_protodict.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 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 = rdf_protodict.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(
rdf_protodict.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 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 = rdf_flows.GrrMessage(request_id=100)
msg_size = len(rdf_protodict.EmbeddedRDFValue(payload=msg)
.SerializeToString())
# Write ~500Kb.
n = 500 * 1024 / msg_size
fd.AddAll([rdf_flows.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 StaticAdd(cls, collection_urn, token, rdf_value, **kwargs):
if not rdf_value.age:
rdf_value.age = rdfvalue.RDFDatetime.Now()
super(GeneralIndexedCollection,
cls).StaticAdd(collection_urn, token,
rdf_protodict.EmbeddedRDFValue(payload=rdf_value),
**kwargs)
def StaticAdd(cls,
collection_urn,
rdf_value,
timestamp=None,
suffix=None,
mutation_pool=None):
if not rdf_value.age:
rdf_value.age = rdfvalue.RDFDatetime.Now()
super(GeneralIndexedCollection,
cls).StaticAdd(collection_urn,
rdf_protodict.EmbeddedRDFValue(payload=rdf_value),
timestamp=timestamp,
suffix=suffix,
mutation_pool=mutation_pool)
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 = rdf_protodict.EmbeddedRDFValue(payload=rdf_value).SerializeToString()
self.fd.Seek(0, 2)
self.fd.Write(struct.pack("<i", len(data)))
self.fd.Write(data)
self.stream_dirty = True
self.size += 1
def StaticAdd(cls, collection_urn, token, rdf_value, **kwargs):
super(GeneralIndexedCollection,
cls).StaticAdd(collection_urn, token,
rdf_protodict.EmbeddedRDFValue(payload=rdf_value),
**kwargs)
def Add(self, rdf_value, **kwargs):
super(GeneralIndexedCollection,
self).Add(rdf_protodict.EmbeddedRDFValue(payload=rdf_value),
**kwargs)
def GenerateSample(self, number=0): return rdf_protodict.EmbeddedRDFValue(rdf_protodict.RDFValueArray([number]))
