def Decorator(cls):
"""Decorator producing 2 classes: legacy style one and a new style one."""
if not legacy_name:
raise ValueError("legacy_name has to be provided")
# Legacy cron jobs have different base classes depending on whether they're
# stateful or not.
if stateful:
aff4_base_cls = StatefulSystemCronFlow
else:
aff4_base_cls = SystemCronFlow
# Make sure that we're dealing with a true mixin to avoid subtle errors.
if issubclass(cls, cronjobs.SystemCronJobBase):
raise ValueError("Mixin class shouldn't inherit from SystemCronJobBase")
if issubclass(cls, aff4_base_cls):
raise ValueError("Mixin class shouldn't inherit from %s" %
aff4_base_cls.__name__)
# Generate legacy class. Register it within the module as it's not going
# to be returned from the decorator.
aff4_cls = compatibility.MakeType(
legacy_name, (cls, LegacyCronJobAdapterMixin, aff4_base_cls), {})
module = sys.modules[cls.__module__]
setattr(module, legacy_name, aff4_cls)
# Generate new class. No need to register it in the module (like the legacy
# one) since it will replace the original decorated class.
reldb_cls = compatibility.MakeType(
compatibility.GetName(cls), (cls, cronjobs.SystemCronJobBase), {})
return reldb_cls
def __init__(cls, name, bases, env_dict): # pylint: disable=no-self-argument
registry.MetaclassRegistry.__init__(cls, name, bases, env_dict)
has_mixin = False
for mixin in itervalues(cls.connection_mixins):
if issubclass(cls, mixin):
has_mixin = True
break
# ApiRegressionTest is a base class, so it doesn't make sense to generate
# _http_v1/_http_v2 classes for it.
# Generating classes from already generated classes would lead to infinite
# recursion. Skipping the generated ones.
if name == "ApiRegressionTest" or has_mixin:
return
for mixin in itervalues(ApiRegressionTestMetaclass.connection_mixins):
if (mixin.skip_legacy_dynamic_proto_tests and
getattr(cls, "uses_legacy_dynamic_protos", False)):
continue
# Do not generate combinations where the mixin demands relational db reads
# but the test is aff4 only.
if (getattr(cls, "aff4_only_test", False) and
getattr(mixin, "read_from_relational_db", False)):
continue
cls_name = "%s_%s" % (name, mixin.connection_type)
test_cls = compatibility.MakeType(
cls_name,
(mixin, cls, test_lib.GRRBaseTest),
# pylint: disable=protected-access
{"testForRegression": lambda x: x._testForRegression()})
module = sys.modules[cls.__module__]
setattr(module, cls_name, test_cls)
def testProtoFileDescriptorIsGeneratedForDynamicAnyValueType(self):
test_pb_file_descriptor, deps = (
DynamicAnyValueTypeTest.EmitProtoFileDescriptor("grr_export"))
pool = descriptor_pool.DescriptorPool()
for file_descriptor in deps + [test_pb_file_descriptor]:
pool.Add(file_descriptor)
proto_descriptor = pool.FindMessageTypeByName(
"grr_export.DynamicAnyValueTypeTest")
factory = message_factory.MessageFactory()
proto_class = factory.GetPrototype(proto_descriptor)
# Now let's define an RDFProtoStruct for the dynamically generated
# proto_class.
new_dynamic_class = compatibility.MakeType(
"DynamicAnyValueTypeTestReversed",
(rdf_structs.RDFProtoStruct, ),
dict(protobuf=proto_class),
)
new_dynamic_instance = new_dynamic_class(type="foo")
self.assertEqual(new_dynamic_instance.type, "foo")
# Test that a proto can be deserialized from serialized RDFValue
# with a dynamic AnyValue field.
test_pb = DynamicAnyValueTypeTest(type="TestStruct")
test_pb.dynamic.foobar = "Hello"
proto_value = proto_class()
proto_value.ParseFromString(test_pb.SerializeToString())
self.assertEqual(proto_value.type, "TestStruct")
self.assertEqual(proto_value.dynamic.type_url, "TestStruct")
self.assertEqual(proto_value.dynamic.value,
test_pb.dynamic.SerializeToString())
def MakeFlatRDFClass(self, value):
"""Generates flattened RDFValue class definition for the given value."""
def Flatten(self, metadata, value_to_flatten):
if metadata:
self.metadata = metadata
for desc in value_to_flatten.type_infos:
if desc.name == "metadata":
continue
if hasattr(self, desc.name) and value_to_flatten.HasField(
desc.name):
setattr(self, desc.name,
getattr(value_to_flatten, desc.name))
descriptors = []
enums = {}
# Metadata is always the first field of exported data.
descriptors.append(
rdf_structs.ProtoEmbedded(name="metadata",
field_number=1,
nested=base.ExportedMetadata))
for number, desc in sorted(value.type_infos_by_field_number.items()):
# Name 'metadata' is reserved to store ExportedMetadata value.
if desc.name == "metadata":
logging.debug("Ignoring 'metadata' field in %s.",
value.__class__.__name__)
continue
# Copy descriptors for primivie values as-is, just make sure their
# field number is correct.
if isinstance(desc,
(rdf_structs.ProtoBinary, rdf_structs.ProtoString,
rdf_structs.ProtoUnsignedInteger,
rdf_structs.ProtoRDFValue, rdf_structs.ProtoEnum)):
# Incrementing field number by 1, as 1 is always occuppied by metadata.
descriptors.append(desc.Copy(field_number=number + 1))
if (isinstance(desc, rdf_structs.ProtoEnum)
and not isinstance(desc, rdf_structs.ProtoBoolean)):
# Attach the enum container to the class for easy reference:
enums[desc.enum_name] = desc.enum_container
# Create the class as late as possible. This will modify a
# metaclass registry, we need to make sure there are no problems.
output_class = compatibility.MakeType(
self.ExportedClassNameForValue(value), (AutoExportedProtoStruct, ),
dict(Flatten=Flatten))
for descriptor in descriptors:
output_class.AddDescriptor(descriptor)
for name, container in enums.items():
setattr(output_class, name, container)
return output_class
def testWithBaseTypes(self):
class Bar(object):
pass
class Baz(object):
pass
cls = compatibility.MakeType("Foo", (Bar, Baz), {})
self.assertEqual(compatibility.GetName(cls), "Foo")
self.assertTrue(issubclass(cls, Bar))
self.assertTrue(issubclass(cls, Baz))
def DualDBTest(cls):
"""Decorator that creates an additional RELDB-enabled test class."""
module = sys.modules[cls.__module__]
cls_name = compatibility.GetName(cls)
db_test_cls_name = "{}_RelationalDBEnabled".format(cls_name)
db_test_cls = compatibility.MakeType(
name=db_test_cls_name,
base_classes=(RelationalDBEnabledMixin, cls),
namespace={})
setattr(module, db_test_cls_name, db_test_cls)
db_test_cls_name = "{}_StableRelationalDBEnabled".format(cls_name)
db_test_cls = compatibility.MakeType(
name=db_test_cls_name,
base_classes=(StableRelationalDBEnabledMixin, cls),
namespace={})
setattr(module, db_test_cls_name, db_test_cls)
return cls
def testWithNamespace(self):
namespace = {
"Bar": lambda self: 42,
"Baz": lambda self, x: x * 2,
}
cls = compatibility.MakeType("Foo", (object,), namespace)
self.assertEqual(compatibility.GetName(cls), "Foo")
foo = cls()
self.assertEqual(foo.Bar(), 42)
self.assertEqual(foo.Baz(42), 84)
def _TestDatabasesDecorator(cls):
"""Decorator that creates additional RELDB-enabled test classes."""
module = sys.modules[cls.__module__]
cls_name = compatibility.GetName(cls)
DualDBTest(cls)
if mysql:
db_test_cls_name = "{}_MySQLEnabled".format(cls_name)
db_test_cls = compatibility.MakeType(
name=db_test_cls_name,
base_classes=(RelationalDBEnabledMixin,
db_test_mixin.GlobalDatabaseTestMixin,
mysql_test.MySQLDatabaseProviderMixin, cls),
namespace={})
setattr(module, db_test_cls_name, db_test_cls)
return cls
def _RDFClass(cls, table: rdf_osquery.OsqueryTable) -> Type[Any]:
"""Creates a dynamic RDF proto struct class for given osquery table.
The fields of the proto will correspond to the columns of the table.
Args:
table: An osquery table for which the class is about to be generated.
Returns:
A class object corresponding to the given table.
"""
rdf_cls_name = "OsqueryTable{}".format(hash(table.query))
try:
return cls._rdf_cls_cache[rdf_cls_name]
except KeyError:
pass
rdf_cls = compatibility.MakeType(rdf_cls_name,
(rdf_structs.RDFProtoStruct, ), {})
rdf_cls.AddDescriptor(
rdf_structs.ProtoEmbedded(name="metadata",
field_number=1,
nested=base.ExportedMetadata))
rdf_cls.AddDescriptor(
rdf_structs.ProtoString(name="__query__", field_number=2))
for idx, column in enumerate(table.header.columns):
# It is possible that RDF column is named "metadata". To avoid name clash
# we must rename it to `__metadata__`.
if column.name == "metadata":
name = "__metadata__"
else:
name = column.name
descriptor = rdf_structs.ProtoString(name=name,
field_number=idx + 3)
rdf_cls.AddDescriptor(descriptor)
cls._rdf_cls_cache[rdf_cls_name] = rdf_cls
return rdf_cls
def _TestDatabasesDecorator(cls):
"""Decorator that creates additional RELDB-enabled test classes."""
module = sys.modules[cls.__module__]
cls_name = compatibility.GetName(cls)
# Prevent MRO issues caused by inheriting the same Mixin multiple times.
base_classes = ()
if not issubclass(cls, db_test_mixin.GlobalDatabaseTestMixin):
base_classes += (db_test_mixin.GlobalDatabaseTestMixin,)
if mysql:
db_test_cls_name = "{}_MySQLEnabled".format(cls_name)
db_test_cls = compatibility.MakeType(
name=db_test_cls_name,
base_classes=base_classes +
(mysql_test.MySQLDatabaseProviderMixin, cls),
namespace={})
setattr(module, db_test_cls_name, db_test_cls)
return cls
def testProtoFileDescriptorIsGeneratedForDynamicType(self):
test_pb_file_descriptor, deps = DynamicTypeTest.EmitProtoFileDescriptor(
"grr_export")
pool = descriptor_pool.DescriptorPool()
for file_descriptor in deps + [test_pb_file_descriptor]:
pool.Add(file_descriptor)
proto_descriptor = pool.FindMessageTypeByName("grr_export.DynamicTypeTest")
factory = message_factory.MessageFactory()
proto_class = factory.GetPrototype(proto_descriptor)
# Now let's define an RDFProtoStruct for the dynamically generated
# proto_class.
new_dynamic_class = compatibility.MakeType(
"DynamicTypeTestReversed", (rdf_structs.RDFProtoStruct,),
dict(protobuf=proto_class, rdf_deps=[rdf_client.User]))
new_dynamic_instance = new_dynamic_class(
type="foo", nested=rdf_client.User(username="******"))
self.assertEqual(new_dynamic_instance.type, "foo")
self.assertEqual(new_dynamic_instance.nested.username, "bar")
def testSimple(self):
cls = compatibility.MakeType("Foo", (object,), {})
self.assertEqual(compatibility.GetName(cls), "Foo")
self.assertIsInstance(cls(), cls)
