def setUp(self): schema_file = os.path.dirname(os.path.realpath(__file__)) + \ "/TCCDMDatum.avsc" self.reader_schema = Utils.load_schema(schema_file) self.writer_schema = Utils.load_schema(schema_file) self.serializer = AvroGenericSerializer(self.writer_schema) self.deserializer = AvroGenericDeserializer(self.reader_schema)
def test_ambiguous_cdm18_parsing(self): deserializer = AvroGenericDeserializer(self.reader_schema, self.reader_schema) fog = record_generator.CDMFileObjectGenerator(self.serializer) ssog = record_generator.CDMSrcSinkObjectGenerator(self.serializer) rkog = record_generator.CDMRegistryKeyObjectGenerator(self.serializer) parser = parsing.CDMParser(self.reader_schema) file_object_record = fog.generate_random_record(5) # Remove all of the unique optional attributes file_object_record["datum"].pop("localPrincipal", None) file_object_record["datum"].pop("size", None) file_object_record["datum"].pop("peInfo", None) file_object_record["datum"].pop("hashes", None) src_sink_object_record = ssog.generate_random_record(5) registry_key_object_record = rkog.generate_random_record(5) sfo = self.serializer.serialize_to_bytes(file_object_record) ssso = self.serializer.serialize_to_bytes(src_sink_object_record) srko = self.serializer.serialize_to_bytes(registry_key_object_record) file_object_record2 = deserializer.deserialize_bytes(sfo) src_sink_object_record2 = deserializer.deserialize_bytes(ssso) registry_key_object_record2 = deserializer.deserialize_bytes(srko) self.assertTrue( parser.get_union_branch_type(file_object_record2) == "FileObject") self.assertTrue( parser.get_union_branch_type(src_sink_object_record2) == "SrcSinkObject") self.assertTrue( parser.get_union_branch_type(registry_key_object_record2) == "RegistryKeyObject")
def test_serialization_nested(self): schema_file = os.path.dirname(os.path.realpath(__file__)) + \ "/LabeledEdge.avsc" schema = Utils.load_schema(schema_file) self.serializer = AvroGenericSerializer(schema) self.deserializer = AvroGenericDeserializer(schema, schema) self.serialization_test_helper(schema, False)
def serialization_test_helper(self, schema, is_union): node_file_path = os.path.dirname(os.path.realpath(__file__)) + \ "/testNodes.avro" edge_file_path = os.path.dirname(os.path.realpath(__file__)) + \ "/testEdges.avro" # Create some nodes and an edge. node1 = Utils.create_node(1, "unitOfExecution", schema, True) node2 = Utils.create_node(2, "artifact", schema, True) edge1 = Utils.create_edge(node1, node2, "read", schema) if is_union: # Serialize the nodes and edge to files. with open(node_file_path, "wb") as node_file: self.serializer = AvroGenericSerializer( self.writer_schema, node_file) self.serializer.serialize_to_file([node1, node2]) self.serializer.close_file_serializer() with open(edge_file_path, "wb") as edge_file: self.serializer = AvroGenericSerializer(self.writer_schema, edge_file) self.serializer.serialize_to_file([edge1]) self.serializer.close_file_serializer() if is_union: # Deserialize from the files to records. with open(node_file_path, "rb") as node_file: self.deserializer = AvroGenericDeserializer( self.reader_schema, self.writer_schema, node_file) deserialized_nodes = \ self.deserializer.deserialize_from_file() self.deserializer.close_file_deserializer() with open(edge_file_path, "rb") as edge_file: self.deserializer = AvroGenericDeserializer( self.reader_schema, self.writer_schema, edge_file) deserialized_edges = \ self.deserializer.deserialize_from_file() self.deserializer.close_file_deserializer() if is_union: # Check the deserialized nodes. self.assertTrue(len(deserialized_nodes) == 2) self.compare_nodes(node1, deserialized_nodes[0]) self.compare_nodes(node2, deserialized_nodes[1]) # Check the deserialized edges. self.assertTrue(len(deserialized_edges) == 1) self.compare_edges(edge1, deserialized_edges[0]) if is_union: # Clean up the files os.remove(node_file_path) os.remove(edge_file_path)
def setUp(self): schema_file = os.path.dirname(os.path.realpath(__file__)) + \ "/LabeledGraph.avsc" schema = Utils.load_schema(schema_file) self.reader_schema = schema self.writer_schema = schema self.node_schema = Utils.get_schema_by_name( self.writer_schema, TestUnionSchema._NODE_SCHEMA_FULLNAME) self.edge_schema = Utils.get_schema_by_name( self.writer_schema, TestUnionSchema._EDGE_SCHEMA_FULLNAME) self.serializer = AvroGenericSerializer(self.writer_schema) self.deserializer = AvroGenericDeserializer(self.reader_schema, self.writer_schema)
def test_serialization_union(self): schema_file = os.path.dirname(os.path.realpath(__file__)) + \ "/LabeledGraphv2.avsc" self.reader_schema = Utils.load_schema(schema_file) self.deserializer = AvroGenericDeserializer(self.reader_schema, self.writer_schema) self.serialization_test_helper(self.writer_schema, True)
class TestUnionSchema(unittest.TestCase): """ Test that serializing and deserializing when using a union schema works as expected. """ _NODE_SCHEMA_FULLNAME = "com.bbn.tc.schema.avro.LabeledNode" _EDGE_SCHEMA_FULLNAME = "com.bbn.tc.schema.avro.LabeledEdge" def setUp(self): schema_file = os.path.dirname(os.path.realpath(__file__)) + \ "/LabeledGraph.avsc" schema = Utils.load_schema(schema_file) self.reader_schema = schema self.writer_schema = schema self.node_schema = Utils.get_schema_by_name( self.writer_schema, TestUnionSchema._NODE_SCHEMA_FULLNAME) self.edge_schema = Utils.get_schema_by_name( self.writer_schema, TestUnionSchema._EDGE_SCHEMA_FULLNAME) self.serializer = AvroGenericSerializer(self.writer_schema) self.deserializer = AvroGenericDeserializer(self.reader_schema, self.writer_schema) def test_union_schema(self): schema = self.writer_schema node1 = Utils.create_node(1, "unitOfExecution", schema, True) node2 = Utils.create_node(2, "agent", schema, True) edge1 = Utils.create_edge(node1, node2, "wasAssociatedWith", schema, True) serialized_node = self.serializer.serialize_to_bytes(node1) serialized_edge = self.serializer.serialize_to_bytes(edge1) deserialized_node = self.deserializer.deserialize_bytes( serialized_node) deserialized_edge = self.deserializer.deserialize_bytes( serialized_edge) # Don't convert these to strings, like in the Java code. That results # in differences due to unicode strings for keys in the # Avro-deserialized in the Python 2.7 version, and we don't really # want to deal with that. self.assertTrue(node1 == deserialized_node) self.assertTrue(edge1 == deserialized_edge)
def setUp(self): union_schema_file = os.path.dirname(os.path.realpath(__file__)) + \ "/LabeledGraph.avsc" self.reader_schema = Utils.load_schema(union_schema_file) self.writer_schema = Utils.load_schema(union_schema_file) self.serializer = AvroGenericSerializer(self.writer_schema) self.deserializer = AvroGenericDeserializer(self.reader_schema, self.writer_schema)
class TestJsonSerialization(TestBase): """ Test of serializing and deserializing records to and from json. """ def setUp(self): schema_file = os.path.dirname(os.path.realpath(__file__)) + \ "/LabeledEdge.avsc" self.schema = Utils.load_schema(schema_file) self.serializer = AvroGenericSerializer(self.schema) self.deserializer = AvroGenericDeserializer(self.schema, self.schema) def test_serialization(self): node1 = Utils.create_node(1, "unitOfExecution", True, self.schema) node2 = Utils.create_node(2, "artifact", True, self.schema) edge = Utils.create_edge(node1, node2, "read", True, self.schema) # Make sure serialization and deserialization is symmetric. json_edge = self.serializer.serialize_to_json(edge, True) deserialized_edge = self.deserializer.deserialize_json(json_edge) self.assertTrue(edge == deserialized_edge) # Make sure that the serializer can serialize to both bytes and json # without corrupting any internal state. Also, test without making # the json serialization prettified. edge = Utils.create_edge(node1, node2, "modified", True, self.schema) self.serializer.serialize_to_bytes(edge) json_edge = self.serializer.serialize_to_json(edge) deserialized_edge = self.deserializer.deserialize_json(json_edge) self.assertTrue(edge == deserialized_edge) def test_bad_record_serialization(self): bad_edge = {"test": "bad"} with self.assertRaises(avro.io.AvroTypeException): self.serializer.serialize_to_json(bad_edge) def test_bad_deserialized_record(self): bad_serialized_edge = '{"test": "bad"}' with self.assertRaises(avro.io.AvroTypeException): self.deserializer.deserialize_json(bad_serialized_edge)
class TestByteSerialization(TestBase): """ Test of serializing to and deserializing from a byte stream. """ def test_serialization_union(self): self.serialization_test_helper(self.writer_schema, True) def test_serialization_nested(self): schema_file = os.path.dirname(os.path.realpath(__file__)) + \ "/LabeledEdge.avsc" schema = Utils.load_schema(schema_file) self.serializer = AvroGenericSerializer(schema) self.deserializer = AvroGenericDeserializer(schema, schema) self.serialization_test_helper(schema, False) def serialization_test_helper(self, schema, is_union): node1 = Utils.create_node(1, "unitOfExecution", schema, True) node2 = Utils.create_node(2, "artifact", schema, True) edge1 = Utils.create_edge(node1, node2, "read", schema) if is_union: serialized_node1 = self.serializer.serialize_to_bytes(node1) deserialized_node1 = \ self.deserializer.deserialize_bytes(serialized_node1) self.compare_nodes(node1, deserialized_node1) serialized_node2 = self.serializer.serialize_to_bytes(node2) deserialized_node2 = \ self.deserializer.deserialize_bytes(serialized_node2) self.compare_nodes(node2, deserialized_node2) serialized_edge1 = self.serializer.serialize_to_bytes(edge1) deserialized_edge1 = \ self.deserializer.deserialize_bytes(serialized_edge1) self.compare_edges(edge1, deserialized_edge1)
def setUp(self): schema_file = os.path.dirname(os.path.realpath(__file__)) + \ "/LabeledEdge.avsc" self.schema = Utils.load_schema(schema_file) self.serializer = AvroGenericSerializer(self.schema) self.deserializer = AvroGenericDeserializer(self.schema, self.schema)
class TestCDMTypeParsing(unittest.TestCase): _KV_PAIRS = 5 def setUp(self): schema_file = os.path.dirname(os.path.realpath(__file__)) + \ "/TCCDMDatum.avsc" self.reader_schema = Utils.load_schema(schema_file) self.writer_schema = Utils.load_schema(schema_file) self.serializer = AvroGenericSerializer(self.writer_schema) self.deserializer = AvroGenericDeserializer(self.reader_schema) def _run_record_type_test(self, generator, expected_value): parser = parsing.CDMParser(self.reader_schema) for i in range(20): record = generator.generate_random_record( TestCDMTypeParsing._KV_PAIRS) self.assertTrue(Utils.validate(self.writer_schema, record)) self.assertTrue(parser.get_record_type(record) == expected_value) serialized = self.serializer.serialize_to_bytes(record) deserialized = self.deserializer.deserialize_bytes(serialized) self.assertTrue( parser.get_record_type(deserialized) == expected_value) def test_provenance_tag_node(self): generator = \ record_generator.CDMProvenanceTagNodeGenerator(self.serializer) self._run_record_type_test(generator, "ProvenanceTagNode") def test_event(self): generator = record_generator.CDMEventGenerator(self.serializer) self._run_record_type_test(generator, "Event") def test_net_flow_object(self): generator = record_generator.CDMNetFlowObjectGenerator(self.serializer) self._run_record_type_test(generator, "NetFlowObject") def test_file_object(self): generator = record_generator.CDMFileObjectGenerator(self.serializer) self._run_record_type_test(generator, "FileObject") def test_src_sink_object(self): generator = record_generator.CDMSrcSinkObjectGenerator(self.serializer) self._run_record_type_test(generator, "SrcSinkObject") def test_ipc_object(self): generator = record_generator.CDMIpcObjectGenerator(self.serializer) self._run_record_type_test(generator, "IpcObject") def test_memory_object(self): generator = record_generator.CDMMemoryObjectGenerator(self.serializer) self._run_record_type_test(generator, "MemoryObject") def test_principal(self): generator = record_generator.CDMPrincipalGenerator(self.serializer) self._run_record_type_test(generator, "Principal") def test_time_marker(self): generator = record_generator.CDMTimeMarkerGenerator(self.serializer) self._run_record_type_test(generator, "TimeMarker") def test_unit_dependency_marker(self): generator = record_generator.CDMUnitDependencyGenerator( self.serializer) self._run_record_type_test(generator, "UnitDependency") def test_registry_key_object(self): generator = record_generator.CDMRegistryKeyObjectGenerator( self.serializer) self._run_record_type_test(generator, "RegistryKeyObject") def test_host(self): generator = record_generator.CDMHostGenerator(self.serializer) self._run_record_type_test(generator, "Host") def test_packet_socket_object(self): generator = record_generator.CDMPacketSocketObjectGenerator( self.serializer) self._run_record_type_test(generator, "PacketSocketObject") def test_end_marker(self): generator = record_generator.CDMEndMarkerGenerator(self.serializer) self._run_record_type_test(generator, "EndMarker") def test_unknown_provenance_node(self): generator = record_generator.CDMUnknownProvenanceNodeGenerator( self.serializer) self._run_record_type_test(generator, "UnknownProvenanceNode")