Merge pull request #270 from avast/feat/remove-variables

feat: Added Rule::removeVariables() to remove all variables from both AST & token stream
avast · Jan 14, 2025 · fd31bb2 · fd31bb2
2 parents bbf9ff0 + ece2e4c
commit fd31bb2
Show file tree

Hide file tree

Showing 4 changed files with 93 additions and 1 deletion.
diff --git a/include/yaramod/types/rule.h b/include/yaramod/types/rule.h
@@ -110,6 +110,7 @@ class Rule
 	void removeString(const std::string& id);
 	void addTag(const std::string& tag);
 	void removeTags(const std::string& tag);
+	void removeVariables();
 	/// @}
 
 private:

diff --git a/src/python/yaramod_python.cpp b/src/python/yaramod_python.cpp
@@ -325,7 +325,8 @@ void addBasicClasses(py::module& module)
 		.def("remove_string", &Rule::removeString)
 		.def("get_meta_with_name", py::overload_cast<const std::string&>(&Rule::getMetaWithName), py::return_value_policy::reference)
 		.def("add_tag", &Rule::addTag)
-		.def("remove_tags", py::overload_cast<const std::string&>(&Rule::removeTags));
+		.def("remove_tags", py::overload_cast<const std::string&>(&Rule::removeTags))
+		.def("remove_variables", &Rule::removeVariables);
 
 	py::class_<Meta>(module, "Meta")
 		.def_property("key", &Meta::getKey, &Meta::setKey)

diff --git a/src/types/rule.cpp b/src/types/rule.cpp
@@ -637,4 +637,24 @@ void Rule::removeTags(TokenType type)
 	}
 }
 
+/**
+ * Removes all variables from a rule.
+ */
+void Rule::removeVariables()
+{
+	auto firstToken = getFirstTokenIt();
+	auto lastToken = getLastTokenIt();
+
+	auto varBegin = _tokenStream->find(TokenType::VARIABLES, firstToken, lastToken);
+	if (varBegin == lastToken)
+		return;
+
+	auto varEnd = _tokenStream->find(TokenType::CONDITION, varBegin, lastToken);
+	if (varEnd == lastToken)
+		return;
+
+	setVariables({});
+	_tokenStream->erase(varBegin, varEnd);
+}
+
 }
diff --git a/tests/cpp/parser_tests.cpp b/tests/cpp/parser_tests.cpp
@@ -8753,5 +8753,75 @@ rule expression_array
 	ASSERT_NE(parent_itr->getElements()[2]->as<BoolLiteralExpression>(), nullptr);
 }
 
+TEST_F(ParserTests,
+RemoveVariablesNone) {
+	prepareInput(
+R"(
+rule remove_variables
+{
+	condition:
+		false
+}
+)");
+
+	EXPECT_TRUE(driver.parse(input, ParserMode::Incomplete));
+	ASSERT_EQ(1u, driver.getParsedFile().getRules().size());
+
+	const auto& rule = driver.getParsedFile().getRules()[0];
+
+	rule->removeVariables();
+
+	EXPECT_EQ(0u, rule->getVariables().size());
+	EXPECT_EQ(R"(rule remove_variables {
+	condition:
+		false
+})", rule->getText());
+
+	EXPECT_EQ(R"(
+rule remove_variables
+{
+	condition:
+		false
+}
+)", driver.getParsedFile().getTextFormatted());
+}
+
+TEST_F(ParserTests,
+RemoveVariables) {
+	prepareInput(
+R"(
+rule remove_variables
+{
+	variables:
+		a = 1
+		b = 2
+		c = a + b
+	condition:
+		false
+}
+)");
+
+	EXPECT_TRUE(driver.parse(input, ParserMode::Incomplete));
+	ASSERT_EQ(1u, driver.getParsedFile().getRules().size());
+
+	const auto& rule = driver.getParsedFile().getRules()[0];
+
+	rule->removeVariables();
+
+	EXPECT_EQ(0u, rule->getVariables().size());
+	EXPECT_EQ(R"(rule remove_variables {
+	condition:
+		false
+})", rule->getText());
+
+	EXPECT_EQ(R"(
+rule remove_variables
+{
+	condition:
+		false
+}
+)", driver.getParsedFile().getTextFormatted());
+}
+
 }
 }