diff --git a/include/yaramod/builder/yara_file_builder.h b/include/yaramod/builder/yara_file_builder.h index 8ddce836..0e223c36 100644 --- a/include/yaramod/builder/yara_file_builder.h +++ b/include/yaramod/builder/yara_file_builder.h @@ -6,6 +6,7 @@ #pragma once +#include #include #include @@ -56,10 +57,13 @@ class YaraFileBuilder YaraFileBuilder& withRule(const std::shared_ptr& rule); /// @} +protected: + void insertImportIntoTokenStream(TokenIt before, const std::string& moduleName); + private: - bool _lastAddedWasImport = false; ///< Flag to determine newlines std::shared_ptr _tokenStream; ///< Tokens storage - std::vector _module_tokens; ///< Modules + std::map _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> _rules; ///< Rules diff --git a/src/builder/yara_file_builder.cpp b/src/builder/yara_file_builder.cpp index c3dd80d6..0f6c0475 100644 --- a/src/builder/yara_file_builder.cpp +++ b/src/builder/yara_file_builder.cpp @@ -20,7 +20,8 @@ namespace yaramod { std::unique_ptr YaraFileBuilder::get(bool recheck, ParserDriver* external_driver) { auto yaraFile = std::make_unique(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 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 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) { - 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) */ YaraFileBuilder& YaraFileBuilder::withRule(const std::shared_ptr& 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; } diff --git a/tests/cpp/builder_tests.cpp b/tests/cpp/builder_tests.cpp index f48ae935..84bda203 100644 --- a/tests/cpp/builder_tests.cpp +++ b/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; diff --git a/tests/python/test_builder.py b/tests/python/test_builder.py index ec8c702a..e3fde73d 100644 --- a/tests/python/test_builder.py +++ b/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):