References for existing structures

include/yaramod/types/modules/module.h

@@ -70,9 +70,12 @@ class Module
 	void _addAttributeFromJson(StructureSymbol* base, const nlohmann::json& json);
 	void _addIterable(StructureSymbol* base, const nlohmann::json& json);
 	void _addFunctions(StructureSymbol* base, const nlohmann::json& json);
-	std::shared_ptr<StructureSymbol> _addStruct(StructureSymbol* base, const nlohmann::json& json);
+	void _addStruct(Symbol* base, const nlohmann::json& json, std::shared_ptr<StructureSymbol>* ref);
+	void _addReference(Symbol* base, const nlohmann::json& json);
 	void _addValue(StructureSymbol* base, const nlohmann::json& json);
 	void _importJson(const nlohmann::json& json);
+	std::shared_ptr<Symbol> _stringToSymbol (const std::shared_ptr<Symbol>& base, const std::string& str);
+	void _addObjectToBase(Symbol* base, std::shared_ptr<Symbol> newAttribute);
 
 	std::string _name; ///< Name of the module
 	std::vector<std::pair<std::string, bool>> _filePaths; ///< The custom paths to JSON files which help to determine this module. Elements: [<path>, true iff <path> was loaded]. May be empty if no private modules.

include/yaramod/types/symbols.h

@@ -37,6 +37,8 @@ class IterableSymbol : public Symbol
 
 	bool isStructured() const { return _elementType == ExpressionType::Object && _structuredType; }
 
+	void setStructuredElementType(std::shared_ptr<Symbol> structure) { _structuredType = std::move(structure); }
+
 protected:
 	IterableSymbol(Symbol::Type type, const std::string& name, ExpressionType elementType, const std::string& documentation)
 		: Symbol(type, name, ExpressionType::Object, documentation), _elementType(elementType), _structuredType() {}
@@ -244,6 +246,12 @@ class StructureSymbol : public Symbol
 		return false;
 	}
 
+	bool addTemplateAttribute(const std::shared_ptr<Symbol>& templateAttribute)
+	{
+		auto insertionResult = _attributes.emplace(templateAttribute->getName(), templateAttribute);
+		return insertionResult.second;
+	}
+
 private:
 	std::unordered_map<std::string, std::shared_ptr<Symbol>> _attributes; ///< Attributes of the structure
 };

src/types/modules/module.cpp

@@ -36,6 +36,34 @@ std::optional<ExpressionType> stringToExpressionType (const std::string& str)
 	return std::nullopt;
 }
 
+/**
+ * Finds an attribute with specified name for a structure or iterable base.
+ */
+std::optional<std::shared_ptr<Symbol>> _getExistingAttribute(Symbol* base, std::string& name)
+{
+	if (base)
+	{
+		switch (base->getType())
+		{
+			case Symbol::Type::Structure:
+				if (auto structure = dynamic_cast<StructureSymbol*>(base))
+					return structure->getAttribute(name);
+				break;
+			case Symbol::Type::Array:
+			case Symbol::Type::Dictionary:
+				if (auto iterable = dynamic_cast<IterableSymbol*>(base))
+				{
+					auto structElement = iterable->getStructuredElementType();
+					return structElement ? std::make_optional(structElement) : std::nullopt;
+				}
+				break;
+			default:
+				break;
+		}
+	}
+	return std::nullopt;
+}
+
 /**
  * Constructor.
  *
@@ -172,21 +200,35 @@ void Module::_addIterable(StructureSymbol* base, const Json& json)
 		if (json.contains("structure"))
 		{
 			auto structureJson = accessJsonSubjson(json, "structure");
-			if (accessJsonString(structureJson, "kind") != "struct")
-				throw ModuleError("Colliding definitions of " + name + " attribute. Expected embedded structure to have kind 'struct'." + getPathsAsString());
+			auto kind = accessJsonString(structureJson, "kind");
+
 			if (accessJsonString(structureJson, "name") != name)
 				throw ModuleError("Colliding definitions of " + name + " attribute. '" + name + "' != '" + accessJsonString(structureJson, "name") + "'." + getPathsAsString());
-			if (structureJson.contains("attributes"))
+
+			if (kind == "reference")
+			{
+				auto type = accessJsonString(structureJson, "type");
+				if (_stringToSymbol(nullptr, type) != existingIterable->getStructuredElementType())
+					throw ModuleError("Colliding definitions of " + name + " attribute. Unxpected referenced type." + getPathsAsString());
+			}
+			else if (kind == "struct")
+			{
+				if (structureJson.contains("attributes"))
+				{
+					if (!existingIterable->isStructured())
+						throw ModuleError("Colliding definitions of " + name + " attribute. Unxpected structured iterable." + getPathsAsString());
+					if (existingIterable->getElementType() != ExpressionType::Object)
+						throw ModuleError("Not object");
+					auto attributes = accessJsonArray(structureJson, "attributes");
+					auto existingEmbeddedStructure = std::static_pointer_cast<StructureSymbol>(existingIterable->getStructuredElementType());
+					for (const auto& attr : attributes)
+						_addAttributeFromJson(existingEmbeddedStructure.get(), attr);
+				}
+			}
+			else
 			{
-				if (!existingIterable->isStructured())
-					throw ModuleError("Colliding definitions of " + name + " attribute. Unxpected structured iterable." + getPathsAsString());
-				if (existingIterable->getElementType() != ExpressionType::Object)
-					throw ModuleError("Not object");
-				auto attributes = accessJsonArray(structureJson, "attributes");
-				auto existingEmbeddedStructure = std::static_pointer_cast<StructureSymbol>(existingIterable->getStructuredElementType());
-				for (const auto& attr : attributes)
-					_addAttributeFromJson(existingEmbeddedStructure.get(), attr);
-			}		
+				throw ModuleError("Colliding definitions of " + name + " attribute. Expected embedded structure to have kind 'struct' or 'reference'." + getPathsAsString());
+			}
 		}
 		else if (existingIterable->isStructured())
 			throw ModuleError("Colliding definitions of " + name + " attribute. Expected structured iterable." + getPathsAsString());
@@ -196,12 +238,20 @@ void Module::_addIterable(StructureSymbol* base, const Json& json)
 		if (json.contains("structure"))
 		{
 			auto structureJson = accessJsonSubjson(json, "structure");
-			auto embeddedStructure = _addStruct(nullptr, structureJson);
 
+			std::shared_ptr<Symbol> templateAttribute;
 			if (isDictionary)
-				base->addAttribute(std::make_shared<DictionarySymbol>(name, embeddedStructure, documentation));
+				templateAttribute = std::make_shared<DictionarySymbol>(name, ExpressionType::Object, documentation);
 			else
-				base->addAttribute(std::make_shared<ArraySymbol>(name, embeddedStructure, documentation));
+				templateAttribute = std::make_shared<ArraySymbol>(name, ExpressionType::Object, documentation);
+
+			base->addTemplateAttribute(templateAttribute);
+
+			if (accessJsonString(structureJson, "kind") == "reference")
+				_addReference(templateAttribute.get(), structureJson);
+			else
+				_addStruct(templateAttribute.get(), structureJson, nullptr);
+
 		}
 		else
 		{
@@ -264,23 +314,24 @@ void Module::_addFunctions(StructureSymbol* base, const Json& json)
 
 /**
  * Creates a structure from supplied json
- * If base is supplied, this method returns nullptr and it either:
- *  - adds structure from json as a attribute of base or
+ * If base is supplied, this method either:
+ *  - adds structure from json as an attribute of base and saves it to ref if supllied or
  *  - it modifies already existing attribute of base with the same name.
- * If base is nullptr, this method returns new Structure constructed from supplied json
+ * If base is nullptr, this method constructs new Structure from supplied json and saves to ref if supllied
  *
+ * @param ref pointer to where the Structure gets saved. Can be nullptr.
  * @param json structure supplied in json to be created ("kind": "struct")
- * @param base already existing Structure which gets the new structure as its attribute. Can be nullptr.
+ * @param base already existing Structure or Iterable which gets the new structure as its attribute. Can be nullptr.
  */
-std::shared_ptr<StructureSymbol> Module::_addStruct(StructureSymbol* base, const Json& json)
+void Module::_addStruct(Symbol* base, const Json& json, std::shared_ptr<StructureSymbol>* ref)
 {
 	assert(accessJsonString(json, "kind") == "struct");
 
 	auto name = accessJsonString(json, "name");
 	auto attributes = accessJsonArray(json, "attributes");
 
-	// Before creating new structure we first look for its existence within base attributes:	
-	std::optional<std::shared_ptr<Symbol>> existing = base ? base->getAttribute(name) : std::nullopt;
+	// Before creating new structure we first look for its existence within base attributes:
+	std::optional<std::shared_ptr<Symbol>> existing = _getExistingAttribute(base, name);
 	if (existing)
 	{
 		if (existing.value()->getType() != Symbol::Type::Structure)
@@ -292,13 +343,57 @@ std::shared_ptr<StructureSymbol> Module::_addStruct(StructureSymbol* base, const
 	else
 	{
 		auto newStructure = std::make_shared<StructureSymbol>(name);
-		for (const auto& attr : attributes)
-			_addAttributeFromJson(newStructure.get(), attr);
-		if (!base)
-			return newStructure;
-		base->addAttribute(newStructure);
+
+		if (ref)
+		{
+			*ref = std::move(newStructure);
+			_addObjectToBase(base, *ref);
+
+			for (const auto& attr : attributes)
+				_addAttributeFromJson(ref->get(), attr);
+		}
+		else
+		{
+			_addObjectToBase(base, newStructure);
+			for (const auto& attr : attributes)
+				_addAttributeFromJson(newStructure.get(), attr);
+		}
+	}
+}
+
+/**
+ * Creates a reference from supplied json and:
+ * adds the new reference as a attribute of structure or iterable base.
+ * When already existing attribute of base with specified name,
+ * this method checks that the references are the same.
+ *
+ * @param json reference supplied in json to be created ("kind": "reference")
+ * @param base already existing Structure or Iterable which gets the new reference as its attribute. Can't be nullptr
+ */
+void Module::_addReference(Symbol* base, const Json& json)
+{
+	assert(accessJsonString(json, "kind") == "reference");
+
+	auto name = accessJsonString(json, "name");
+	auto symbol = _stringToSymbol(nullptr, accessJsonString(json, "type"));
+
+	if (!symbol)
+		throw ModuleError("Unknown symbol '" + accessJsonString(json, "type") + "'");
+
+	// Before creating new reference we first look for its existence within base attributes:
+	std::optional<std::shared_ptr<Symbol>> existing = _getExistingAttribute(base, name);
+	if (existing)
+	{
+		if (existing.value()->getType() != Symbol::Type::Reference)
+			throw ModuleError("Colliding definitions of " + name + " attribute with different kind. Expected reference." + getPathsAsString());
+		auto existingReference = std::static_pointer_cast<ReferenceSymbol>(existing.value());
+		if (existingReference->getSymbol() != symbol)
+			throw ModuleError("Colliding definitions of " + name + " attribute. The value is defined twice with different references. " + getPathsAsString());
+	}
+	else
+	{
+		_addObjectToBase(base, std::make_shared<ReferenceSymbol>(name, symbol));
 	}
-	return nullptr;
 }
 
 /**
@@ -351,7 +446,9 @@ void Module::_addAttributeFromJson(StructureSymbol* base, const Json& json)
 	if (kind == "function")
 		_addFunctions(base, json);
 	else if (kind == "struct")
-		_addStruct(base, json);
+		_addStruct(base, json, nullptr);
+	else if (kind == "reference")
+		_addReference(base, json);
 	else if (kind == "value")
 		_addValue(base, json);
 	else if (kind == "dictionary" || kind == "array")
@@ -392,7 +489,7 @@ void Module::_importJson(const Json& json)
 	if (name == std::string{})
 		throw ModuleError("Module name must be non-empty.");
 	else if (!_structure) // First iteration - need to create the structure.
-		_structure = _addStruct(nullptr, json);
+		_addStruct(nullptr, json, &_structure);
 	else if (_name != name) // Throws - name of the module must be the same accross the files.
 		throw ModuleError("Module name must be the same in all files, but " + name + " != " + _name + ".\n" + getPathsAsString());
 	else // _struct already created, need only to add new attributes
@@ -403,4 +500,87 @@ void Module::_importJson(const Json& json)
 	}
 }
 
+/**
+ * A mapping converting a given string to corresponding Symbol.
+ */
+std::shared_ptr<Symbol> Module::_stringToSymbol (const std::shared_ptr<Symbol>& base, const std::string& str)
+{
+	std::string current_attribute = str;
+	std::string next_attribute;
+	auto delim_index = str.find('.');
+
+	if (delim_index != std::string::npos)
+	{
+		current_attribute = str.substr(0, delim_index);
+		next_attribute = str.substr(delim_index + 1, str.length());
+	}
+
+	if (!base)
+	{
+		if (_structure->getName() != current_attribute)
+			throw ModuleError("Unaccessible namespace. You can reference objects only from the same module." + getPathsAsString());
+
+		if (!next_attribute.empty())
+			return _stringToSymbol(_structure, next_attribute);
+		else
+			return _structure;
+	}
+	else
+	{
+		if (base->getType() != Symbol::Type::Structure)
+			throw ModuleError("Can't access an attribute of " + base->getName() + " as it's not a structured object." + getPathsAsString());
+
+		auto baseStruct = std::static_pointer_cast<StructureSymbol>(base);
+		auto found_symbol = baseStruct->getAttribute(current_attribute);
+
+		if (!found_symbol)
+			throw ModuleError("Object " + base->getName() + " does not contain an attribute with the name of " + current_attribute + "." + getPathsAsString());
+
+		auto current_symbol = found_symbol.value();
+
+		if (current_symbol->getType() == Symbol::Type::Dictionary ||
+			current_symbol->getType() == Symbol::Type::Array)
+		{
+			auto current_iterable = std::static_pointer_cast<IterableSymbol>(current_symbol);
+			current_symbol = current_iterable->getStructuredElementType();
+		}
+
+		if (!next_attribute.empty())
+		{
+			auto newBase = std::static_pointer_cast<StructureSymbol>(current_symbol);
+			return _stringToSymbol(newBase, next_attribute);
+		}
+
+		return current_symbol;
+	}
+}
+
+/**
+ * Adds an attribute to it's base that can be either a parent structure or an iterable.
+ */
+void Module::_addObjectToBase(Symbol* base, std::shared_ptr<Symbol> newAttribute)
+{
+	if (base)
+	{
+		switch (base->getType())
+		{
+			case Symbol::Type::Structure:
+				if (auto structure = dynamic_cast<StructureSymbol*>(base))
+					structure->addAttribute(newAttribute);
+				else
+					throw ModuleError("Base could not be casted to a structure." + getPathsAsString());
+				break;
+			case Symbol::Type::Array:
+			case Symbol::Type::Dictionary:
+				if (auto iterable = dynamic_cast<IterableSymbol*>(base))
+					iterable->setStructuredElementType(newAttribute);
+				else
+					throw ModuleError("Base could not be casted to an iterable." + getPathsAsString());
+				break;
+			default:
+				throw ModuleError("Base type has to be a structure or an iterable." + getPathsAsString());
+		}
+	}
+}
+
 }

tests/python/test_builder.py

@@ -1290,6 +1290,7 @@ def test_rule_with_string_modifiers(self):
     def test_rule_with_custom_modules(self):
         cond = yaramod.conjunction([
             yaramod.id("module_test.structure_test.function_test")(yaramod.regexp("abc", "")),
+            yaramod.id("module_test.reference_test.function_test")(yaramod.regexp("cba", "")),
             yaramod.id("cuckoo.sync.mutex")(yaramod.regexp("abc", ""))
         ]).get()
         rule = yaramod.YaraRuleBuilder() \
@@ -1309,6 +1310,7 @@ def test_rule_with_custom_modules(self):
 {
 	condition:
 		module_test.structure_test.function_test(/abc/) and
+		module_test.reference_test.function_test(/cba/) and
 		cuckoo.sync.mutex(/abc/)
 }
 ''')
@@ -1317,7 +1319,7 @@ def test_rule_with_custom_modules(self):
 
 rule test {
 	condition:
-		module_test.structure_test.function_test(/abc/) and cuckoo.sync.mutex(/abc/)
+		module_test.structure_test.function_test(/abc/) and module_test.reference_test.function_test(/cba/) and cuckoo.sync.mutex(/abc/)
 }''')
 
     def test_rule_with_defined_condition(self):

tests/python/test_representation.py

@@ -295,6 +295,16 @@ def test_custom_module_interface(self):
         self.assertTrue(value_symbol.is_value)
         self.assertEqual(value_symbol.documentation, "Testing value documentation. Example: ```module_test.value_test > 10```")
 
+        self.assertTrue("reference_test" in cuckoo_attributes)
+        reference_symbol = cuckoo_attributes["reference_test"]
+        self.assertTrue(reference_symbol.is_reference)
+        self.assertEqual(reference_symbol.symbol, structure_symbol)
+
+        self.assertTrue("references_test" in cuckoo_attributes)
+        references_symbol = cuckoo_attributes["references_test"]
+        self.assertTrue(references_symbol.is_array)
+        self.assertTrue(references_symbol.structure.is_reference)
+        self.assertEqual(references_symbol.structure.symbol, structure_symbol)
 
     def test_custom_module_enhancing_known_module(self):
         modules = yaramod.Yaramod(yaramod.Features.AllCurrent, "./tests/python/testing_modules").modules