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)