class StructGrrMessage(rdf_structs.RDFProtoStruct):
"""A serialization agnostic GrrMessage."""
type_description = type_info.TypeDescriptorSet(
rdf_structs.ProtoString(
name="session_id",
field_number=1,
description="Every Flow has a unique session id."),
rdf_structs.ProtoUnsignedInteger(
name="request_id",
field_number=2,
description="This message is in response to this request number"),
rdf_structs.ProtoUnsignedInteger(
name="response_id",
field_number=3,
description="Responses for each request."),
rdf_structs.ProtoString(
name="name",
field_number=4,
description=("This is the name of the client action that will be "
"executed. It is set by the flow and is executed by "
"the client.")),
rdf_structs.ProtoBinary(
name="args",
field_number=5,
description="This field contains an encoded rdfvalue."),
rdf_structs.ProtoString(
name="source",
field_number=6,
description=("Client name where the message came from (This is "
"copied from the MessageList)")),)
class FastGrrMessageList(rdf_structs.RDFProtoStruct):
"""A Faster implementation of GrrMessageList."""
type_description = type_info.TypeDescriptorSet(
rdf_structs.ProtoList(
rdf_structs.ProtoEmbedded(
name="job", field_number=1, nested=StructGrrMessage)))
class SqliteTestStruct(rdf_structs.RDFProtoStruct):
"""Custom struct for testing schema generation."""
type_description = type_info.TypeDescriptorSet(
rdf_structs.ProtoString(name="string_field", field_number=1),
rdf_structs.ProtoBinary(name="bytes_field", field_number=2),
rdf_structs.ProtoUnsignedInteger(name="uint_field", field_number=3),
rdf_structs.ProtoSignedInteger(name="int_field", field_number=4),
rdf_structs.ProtoFloat(name="float_field", field_number=5),
rdf_structs.ProtoDouble(name="double_field", field_number=6),
rdf_structs.ProtoEnum(name="enum_field",
field_number=7,
enum_name="EnumField",
enum={
"FIRST": 1,
"SECOND": 2
}),
rdf_structs.ProtoBoolean(name="bool_field", field_number=8),
rdf_structs.ProtoRDFValue(name="urn_field",
field_number=9,
rdf_type="RDFURN"),
rdf_structs.ProtoRDFValue(name="time_field",
field_number=10,
rdf_type="RDFDatetime"),
rdf_structs.ProtoRDFValue(name="time_field_seconds",
field_number=11,
rdf_type="RDFDatetimeSeconds"),
rdf_structs.ProtoRDFValue(name="duration_field",
field_number=12,
rdf_type="Duration"),
rdf_structs.ProtoEmbedded(name="embedded_field",
field_number=13,
nested=TestEmbeddedStruct))
def testTypeDescriptorSet(self):
type_infos = [
type_info.String(name="output", default="analysis/{p}/{u}-{t}"),
type_info.String(description="Profile to use.",
name="profile",
default=""),
type_info.String(description="A comma separated list of plugins.",
name="plugins",
default=""),
]
info = type_info.TypeDescriptorSet(
type_infos[0],
type_infos[1],
type_infos[2],
)
new_info = type_info.TypeDescriptorSet(type_infos[0], )
updated_info = new_info + type_info.TypeDescriptorSet(type_infos[1], )
updated_info += type_info.TypeDescriptorSet(type_infos[2], )
self.assertEqual(info.descriptor_map, updated_info.descriptor_map)
self.assertEqual(sorted(info.descriptors),
sorted(updated_info.descriptors))
self.assertTrue(type_infos[1] in updated_info.descriptors)
self.assertTrue("plugins" in updated_info)
removed_info = updated_info.Remove("plugins")
self.assertTrue(type_infos[1] in updated_info.descriptors)
self.assertTrue("plugins" in updated_info)
self.assertFalse(type_infos[2] in removed_info.descriptors)
self.assertFalse("plugins" in removed_info)
class DynamicAnyValueTypeTest(rdf_structs.RDFProtoStruct):
"""A protobuf with dynamic types stored in AnyValue messages."""
type_description = type_info.TypeDescriptorSet(
rdf_structs.ProtoString(
name="type",
field_number=1,
# By default return the TestStruct proto.
description="A string value"),
rdf_structs.ProtoDynamicAnyValueEmbedded(
name="dynamic",
# The callback here returns the type specified by the type member.
dynamic_cb=lambda x: rdf_structs.RDFProtoStruct.classes.get(x.type
),
field_number=2,
description="A dynamic value based on another field."),
)
class LateBindingTest(rdf_structs.RDFProtoStruct):
type_description = type_info.TypeDescriptorSet(
# A nested protobuf referring to an undefined type.
rdf_structs.ProtoEmbedded(name="nested",
field_number=1,
nested="UndefinedYet"),
rdf_structs.ProtoRDFValue(name="rdfvalue",
field_number=6,
rdf_type="UndefinedRDFValue",
description="An undefined RDFValue field."),
# A repeated late bound field.
rdf_structs.ProtoList(
rdf_structs.ProtoRDFValue(
name="repeated",
field_number=7,
rdf_type="UndefinedRDFValue2",
description="An undefined RDFValue field.")),
)
class UnionTest(rdf_structs.RDFProtoStruct):
union_field = "struct_flavor"
type_description = type_info.TypeDescriptorSet(
rdf_structs.ProtoEnum(name="struct_flavor",
field_number=1,
enum_name="Type",
enum=dict(FIRST=1, SECOND=2, THIRD=3),
default=3,
description="An union enum field"),
rdf_structs.ProtoFloat(name="first",
field_number=2,
description="A float number",
default=1.1),
rdf_structs.ProtoString(name="second",
field_number=3,
default="string",
description="A string value"),
rdf_structs.ProtoUnsignedInteger(name="third",
field_number=4,
default=5,
description="An integer value"),
)
class TestStruct(rdf_structs.RDFProtoStruct):
"""A test struct object."""
type_description = type_info.TypeDescriptorSet(
rdf_structs.ProtoString(
name="foobar",
field_number=1,
default="string",
description="A string value",
labels=[rdf_structs.SemanticDescriptor.Labels.HIDDEN]),
rdf_structs.ProtoUnsignedInteger(name="int",
field_number=2,
default=5,
description="An integer value"),
rdf_structs.ProtoList(
rdf_structs.ProtoString(name="repeated",
field_number=3,
description="A repeated string value")),
# We can serialize an arbitrary RDFValue. This will be serialized into a
# binary string and parsed on demand.
rdf_structs.ProtoRDFValue(name="urn",
field_number=6,
default=rdfvalue.RDFURN("www.google.com"),
rdf_type="RDFURN",
description="An arbitrary RDFValue field."),
rdf_structs.ProtoEnum(name="type",
field_number=7,
enum_name="Type",
enum=dict(FIRST=1, SECOND=2, THIRD=3),
default=3,
description="An enum field"),
rdf_structs.ProtoFloat(name="float",
field_number=8,
description="A float number",
default=1.1),
)
class PartialTest1(rdf_structs.RDFProtoStruct):
"""This is a protobuf with fewer fields than TestStruct."""
type_description = type_info.TypeDescriptorSet(
rdf_structs.ProtoUnsignedInteger(name="int", field_number=2), )
class UndefinedYet(rdf_structs.RDFProtoStruct):
type_description = type_info.TypeDescriptorSet(
rdf_structs.ProtoString(name="foobar",
field_number=1,
description="A string value"), )
class TestEmbeddedStruct(rdf_structs.RDFProtoStruct):
"""Custom struct for testing schema generation."""
type_description = type_info.TypeDescriptorSet(
rdf_structs.ProtoString(name="e_string_field", field_number=1),
rdf_structs.ProtoDouble(name="e_double_field", field_number=2))
