Merge pull request #78 from avast/fix_builder_imports

Builder: avoid duplicity of imports, sort imports lexicographically

include/yaramod/builder/yara_file_builder.h

@@ -6,6 +6,7 @@
 
 #pragma once
 
+#include <map>
 #include <memory>
 #include <vector>
 
@@ -56,10 +57,13 @@ class YaraFileBuilder
 	YaraFileBuilder& withRule(const std::shared_ptr<Rule>& rule);
 	/// @}
 
+protected:
+	void insertImportIntoTokenStream(TokenIt before, const std::string& moduleName);
+
 private:
-	bool _lastAddedWasImport = false; ///< Flag to determine newlines
 	std::shared_ptr<TokenStream> _tokenStream; ///< Tokens storage
-	std::vector<TokenIt> _module_tokens; ///< Modules
+	std::map<std::string, TokenIt> _module_tokens; ///< Modules
+	TokenIt _newline_after_imports; ///< Always stands behind newline after last import in the TokenStream
 	ImportFeatures _import_features; ///< Determines which modules should be possible to load
 	ModulesPool _modules_pool; ///< Storage of used modules
 	std::vector<std::shared_ptr<Rule>> _rules; ///< Rules

src/builder/yara_file_builder.cpp

@@ -20,7 +20,8 @@ namespace yaramod {
 std::unique_ptr<YaraFile> YaraFileBuilder::get(bool recheck, ParserDriver* external_driver)
 {
 	auto yaraFile = std::make_unique<YaraFile>(std::move(_tokenStream), _import_features);
-	yaraFile->addImports(_module_tokens, _modules_pool);
+	for (const auto& module_token : _module_tokens)
+		yaraFile->addImport(module_token.second, _modules_pool);
 	yaraFile->addRules(_rules);
 
 	_module_tokens.clear();
@@ -67,6 +68,14 @@ std::unique_ptr<YaraFile> YaraFileBuilder::get(bool recheck, ParserDriver* exter
 	return yaraFile;
 }
 
+void YaraFileBuilder::insertImportIntoTokenStream(TokenIt before, const std::string& moduleName)
+{
+	_tokenStream->emplace(before, TokenType::IMPORT_KEYWORD, "import");
+	auto moduleToken = _tokenStream->emplace(before, TokenType::IMPORT_MODULE, moduleName);
+	_tokenStream->emplace(before, NEW_LINE, "\n");
+	_module_tokens.insert(std::make_pair(moduleName, moduleToken));
+}
+
 /**
  * Adds module to YARA file.
  *
@@ -76,16 +85,25 @@ std::unique_ptr<YaraFile> YaraFileBuilder::get(bool recheck, ParserDriver* exter
  */
 YaraFileBuilder& YaraFileBuilder::withModule(const std::string& moduleName)
 {
-	if (!_module_tokens.empty())
+	TokenIt moduleToken;
+	if (_module_tokens.empty())
 	{
-		if (!_lastAddedWasImport)
-			_tokenStream->emplace_back(NEW_LINE, "\n");
+		auto before = _tokenStream->begin();
+		insertImportIntoTokenStream(before, moduleName);
+		_newline_after_imports = _tokenStream->emplace(before, NEW_LINE, "\n");
 	}
-	_tokenStream->emplace_back(TokenType::IMPORT_KEYWORD, "import");
-	TokenIt moduleToken = _tokenStream->emplace_back(TokenType::IMPORT_MODULE, moduleName);
-	_tokenStream->emplace_back(NEW_LINE, "\n");
-	_module_tokens.push_back(moduleToken);
-	_lastAddedWasImport = true;
+	else
+	{
+		auto close_module = _module_tokens.lower_bound(moduleName);
+		if (close_module == _module_tokens.end())
+ 			insertImportIntoTokenStream(_newline_after_imports, moduleName);
+		else if (close_module->first != moduleName)
+		{
+			auto before = _tokenStream->findBackwards(TokenType::IMPORT_KEYWORD, close_module->second);
+			insertImportIntoTokenStream(before, moduleName);
+		}
+	}
+
 	return *this;
 }
 
@@ -111,14 +129,13 @@ YaraFileBuilder& YaraFileBuilder::withRule(Rule&& rule)
  */
 YaraFileBuilder& YaraFileBuilder::withRule(std::unique_ptr<Rule>&& rule)
 {
-	if (!_rules.empty() || _lastAddedWasImport)
+	if (!_rules.empty())
 		_tokenStream->emplace_back(NEW_LINE, "\n");
 
 	_tokenStream->move_append(rule->getTokenStream());
 	_tokenStream->emplace_back(NEW_LINE, "\n");
 
 	_rules.emplace_back(std::move(rule));
-	_lastAddedWasImport = false;
 	return *this;
 }
 
@@ -131,13 +148,12 @@ YaraFileBuilder& YaraFileBuilder::withRule(std::unique_ptr<Rule>&& rule)
  */
 YaraFileBuilder& YaraFileBuilder::withRule(const std::shared_ptr<Rule>& rule)
 {
-	if (!_rules.empty() || _lastAddedWasImport)
+	if (!_rules.empty() || !_module_tokens.empty())
 		_tokenStream->emplace_back(NEW_LINE, "\n");
 
 	_tokenStream->move_append(rule->getTokenStream());
 	_tokenStream->emplace_back(NEW_LINE, "\n");
 	_rules.emplace_back(rule);
-	_lastAddedWasImport = false;
 	return *this;
 }
 

tests/cpp/builder_tests.cpp

@@ -41,8 +41,36 @@ PureImportsWorks) {
 		.get(true, &driver);
 
 	ASSERT_NE(nullptr, yaraFile);
-	EXPECT_EQ(R"(import "pe"
+	EXPECT_EQ(R"(import "elf"
+import "pe"
+)", yaraFile->getText());
+}
+
+TEST_F(BuilderTests,
+PureImportsComplicateWorks) {
+	YaraFileBuilder newFile;
+	auto yaraFile = newFile
+		.withModule("pe")
+		.withModule("cuckoo")
+		.withModule("elf")
+		.get(true, &driver);
+
+	ASSERT_NE(nullptr, yaraFile);
+	EXPECT_EQ(R"(import "cuckoo"
 import "elf"
+import "pe"
+)", yaraFile->getText());
+
+	yaraFile = newFile
+		.withModule("cuckoo")
+		.withModule("pe")
+		.withModule("elf")
+		.get(true, &driver);
+
+	ASSERT_NE(nullptr, yaraFile);
+	EXPECT_EQ(R"(import "cuckoo"
+import "elf"
+import "pe"
 )", yaraFile->getText());
 }
 
@@ -75,6 +103,78 @@ UnnamedRuleWorks) {
 )", yaraFile->getTextFormatted());
 }
 
+TEST_F(BuilderTests,
+UnnamedRuleWithImportsWorks) {
+	YaraRuleBuilder newRule;
+	auto rule = newRule.get();
+
+	YaraFileBuilder newFile;
+	auto yaraFile = newFile
+		.withRule(std::move(rule))
+		.withModule("cuckoo")
+		.withModule("pe")
+		.withModule("cuckoo")
+		.withModule("elf")
+		.get(true);
+
+	ASSERT_NE(nullptr, yaraFile);
+	EXPECT_EQ(R"(import "cuckoo"
+import "elf"
+import "pe"
+
+rule unknown {
+	condition:
+		true
+})", yaraFile->getText());
+
+	EXPECT_EQ(R"(import "cuckoo"
+import "elf"
+import "pe"
+
+rule unknown
+{
+	condition:
+		true
+}
+)", yaraFile->getTextFormatted());
+}
+
+TEST_F(BuilderTests,
+UnnamedRuleWithImportsWorks2) {
+	YaraRuleBuilder newRule;
+	auto rule = newRule.get();
+
+	YaraFileBuilder newFile;
+	auto yaraFile = newFile
+		.withModule("cuckoo")
+		.withModule("pe")
+		.withRule(std::move(rule))
+		.withModule("cuckoo")
+		.withModule("elf")
+		.get(true);
+
+	ASSERT_NE(nullptr, yaraFile);
+	EXPECT_EQ(R"(import "cuckoo"
+import "elf"
+import "pe"
+
+rule unknown {
+	condition:
+		true
+})", yaraFile->getText());
+
+	EXPECT_EQ(R"(import "cuckoo"
+import "elf"
+import "pe"
+
+rule unknown
+{
+	condition:
+		true
+}
+)", yaraFile->getTextFormatted());
+}
+
 TEST_F(BuilderTests,
 RuleWithCustomNameWorks) {
 	YaraRuleBuilder newRule;

tests/python/test_builder.py

@@ -19,13 +19,14 @@ def test_pure_imports(self):
             .with_module('elf') \
             .get()
 
-        self.assertEqual(yara_file.text_formatted, '''import "phish"
+        self.assertEqual(yara_file.text_formatted, '''import "elf"
 import "pe"
-import "elf"
+import "phish"
+
 ''')
-        self.assertEqual(yara_file.text, '''import "phish"
+        self.assertEqual(yara_file.text, '''import "elf"
 import "pe"
-import "elf"
+import "phish"
 ''')
 
     def test_empty_rule(self):