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]))