Add unary operator defined

avast · Jul 16, 2021 · 3acd43a · 3acd43a
1 parent 14b5b84
commit 3acd43a
Show file tree

Hide file tree

Showing 14 changed files with 141 additions and 1 deletion.
diff --git a/include/yaramod/builder/yara_expression_builder.h b/include/yaramod/builder/yara_expression_builder.h
@@ -185,6 +185,7 @@ class YaraExpressionBuilder
 
 	YaraExpressionBuilder& contains(const YaraExpressionBuilder& other);
 	YaraExpressionBuilder& matches(const YaraExpressionBuilder& other);
+	YaraExpressionBuilder &defined();
 
 	YaraExpressionBuilder& access(const std::string& attr);
 	YaraExpressionBuilder& operator[](const YaraExpressionBuilder& other);

diff --git a/include/yaramod/types/expressions.h b/include/yaramod/types/expressions.h
@@ -387,7 +387,7 @@ class UnaryOpExpression : public Expression
 public:
 	virtual std::string getText(const std::string& indent = std::string{}) const override
 	{
-		if (_op->getType() == TokenType::NOT)
+		if (_op->getType() == TokenType::NOT || _op->getType() == TokenType::DEFINED)
 			return _op->getString() + " " + _expr->getText(indent);
 		else
 			return _op->getString() + _expr->getText(indent);
@@ -439,6 +439,26 @@ class NotExpression : public UnaryOpExpression
 	}
 };
 
+/**
+ * Class representing defined operation.
+ *
+ * For example:
+ * @code
+ * defined @str
+ * @endcode
+ */
+class DefinedExpression : public UnaryOpExpression
+{
+public:
+	template<typename ExpPtr>
+	DefinedExpression(TokenIt op, ExpPtr &&expr) : UnaryOpExpression(op, std::forward<ExpPtr>(expr)) {}
+
+	virtual VisitResult accept(Visitor *v) override
+	{
+		return v->visit(this);
+	}
+};
+
 /**
  * Class representing unary minus operation.
  *

diff --git a/include/yaramod/types/token_type.h b/include/yaramod/types/token_type.h
@@ -80,6 +80,7 @@ enum class TokenType
 	IMPORT_MODULE,
 	IMPORT_KEYWORD,
 	NOT,
+	DEFINED,
 	AND,
 	OR,
 	ALL,

diff --git a/include/yaramod/utils/modifying_visitor.h b/include/yaramod/utils/modifying_visitor.h
@@ -150,6 +150,11 @@ class ModifyingVisitor : public Visitor
 		return _handleUnaryOperation(expr);
 	}
 
+	virtual VisitResult visit(DefinedExpression *expr) override
+	{
+		return _handleUnaryOperation(expr);
+	}
+
 	virtual VisitResult visit(UnaryMinusExpression* expr) override
 	{
 		return _handleUnaryOperation(expr);

diff --git a/include/yaramod/utils/observing_visitor.h b/include/yaramod/utils/observing_visitor.h
@@ -72,6 +72,12 @@ class ObservingVisitor : public Visitor
 		return {};
 	}
 
+	virtual VisitResult visit(DefinedExpression *expr) override
+	{
+		expr->getOperand()->accept(this);
+		return {};
+	}
+
 	virtual VisitResult visit(UnaryMinusExpression* expr) override
 	{
 		expr->getOperand()->accept(this);

diff --git a/include/yaramod/utils/visitor.h b/include/yaramod/utils/visitor.h
@@ -20,6 +20,7 @@ class StringCountExpression;
 class StringOffsetExpression;
 class StringLengthExpression;
 class NotExpression;
+class DefinedExpression;
 class UnaryMinusExpression;
 class BitwiseNotExpression;
 class AndExpression;
@@ -84,6 +85,7 @@ class Visitor
 	virtual VisitResult visit(StringOffsetExpression* expr) = 0;
 	virtual VisitResult visit(StringLengthExpression* expr) = 0;
 	virtual VisitResult visit(NotExpression* expr) = 0;
+	virtual VisitResult visit(DefinedExpression* expr) = 0;
 	virtual VisitResult visit(UnaryMinusExpression* expr) = 0;
 	virtual VisitResult visit(BitwiseNotExpression* expr) = 0;
 	virtual VisitResult visit(AndExpression* expr) = 0;

diff --git a/src/builder/yara_expression_builder.cpp b/src/builder/yara_expression_builder.cpp
@@ -529,6 +529,18 @@ YaraExpressionBuilder& YaraExpressionBuilder::matches(const YaraExpressionBuilde
 	return *this;
 }
 
+/**
+ * Applies operation defined on two expression.
+ *
+ * @return Builder.
+ */
+YaraExpressionBuilder &YaraExpressionBuilder::defined() {
+	TokenIt token = _tokenStream->emplace(_tokenStream->begin(), TokenType::DEFINED, "defined");
+	_expr = std::make_shared<DefinedExpression>(token, std::move(_expr));
+	setType(Expression::Type::Int);
+	return *this;
+}
+
 /**
  * Accesses the attribute of a structure expression.
  *

diff --git a/src/parser/parser_driver.cpp b/src/parser/parser_driver.cpp
@@ -134,6 +134,13 @@ void ParserDriver::defineTokens()
 	_parser.token("import").symbol("IMPORT_KEYWORD").description("import").action([&](std::string_view str) -> Value { return emplace_back(TokenType::IMPORT_KEYWORD, std::string{str}); });
 	_parser.token("not").symbol("NOT").description("not").action([&](std::string_view str) -> Value { return emplace_back(TokenType::NOT, std::string{str}); })
 		.precedence(14, pog::Associativity::Right);
+	if (_features & Features::AvastOnly) {
+		_parser.token("defined").symbol("DEFINED").description("defined").action(
+						[&](std::string_view str) -> Value {
+							return emplace_back(TokenType::DEFINED, std::string{str});
+						})
+				.precedence(15, pog::Associativity::Right);
+	}
 	_parser.token("and").symbol("AND").description("and").action([&](std::string_view str) -> Value { return emplace_back(TokenType::AND, std::string{str}); })
 		.precedence(5, pog::Associativity::Left);
 	_parser.token("or").symbol("OR").description("or").action([&](std::string_view str) -> Value { return emplace_back(TokenType::OR, std::string{str}); })
@@ -1227,6 +1234,14 @@ void ParserDriver::defineGrammar()
 			output->setTokenStream(currentTokenStream());
 			return output;
 		})
+		.production("DEFINED", "expression", [&](auto&& args) -> Value {
+			TokenIt not_token = args[0].getTokenIt();
+			auto expr = std::move(args[1].getExpression());
+			auto output = std::make_shared<DefinedExpression>(not_token, std::move(expr));
+			output->setType(Expression::Type::Bool);
+			output->setTokenStream(currentTokenStream());
+			return output;
+			})
 		.production("expression", "AND", "expression", [&](auto&& args) -> Value {
 			auto left = std::move(args[0].getExpression());
 			TokenIt and_token = args[1].getTokenIt();

diff --git a/src/python/py_visitor.cpp b/src/python/py_visitor.cpp
@@ -27,6 +27,7 @@ void addVisitorClasses(py::module& module)
 		.def("visit_StringOffsetExpression", py::overload_cast<StringOffsetExpression*>(&Visitor::visit))
 		.def("visit_StringLengthExpression", py::overload_cast<StringLengthExpression*>(&Visitor::visit))
 		.def("visit_NotExpression", py::overload_cast<NotExpression*>(&Visitor::visit))
+		.def("visit_DefinedExpression", py::overload_cast<DefinedExpression*>(&Visitor::visit))
 		.def("visit_UnaryMinusExpression", py::overload_cast<UnaryMinusExpression*>(&Visitor::visit))
 		.def("visit_BitwiseNotExpression", py::overload_cast<BitwiseNotExpression*>(&Visitor::visit))
 		.def("visit_AndExpression", py::overload_cast<AndExpression*>(&Visitor::visit))
@@ -84,6 +85,7 @@ void addVisitorClasses(py::module& module)
 		.def("visit_StringOffsetExpression", py::overload_cast<StringOffsetExpression*>(&ObservingVisitor::visit))
 		.def("visit_StringLengthExpression", py::overload_cast<StringLengthExpression*>(&ObservingVisitor::visit))
 		.def("visit_NotExpression", py::overload_cast<NotExpression*>(&ObservingVisitor::visit))
+		.def("visit_DefinedExpression", py::overload_cast<DefinedExpression*>(&ObservingVisitor::visit))
 		.def("visit_UnaryMinusExpression", py::overload_cast<UnaryMinusExpression*>(&ObservingVisitor::visit))
 		.def("visit_BitwiseNotExpression", py::overload_cast<BitwiseNotExpression*>(&ObservingVisitor::visit))
 		.def("visit_AndExpression", py::overload_cast<AndExpression*>(&ObservingVisitor::visit))
@@ -142,6 +144,7 @@ void addVisitorClasses(py::module& module)
 		.def("visit_StringOffsetExpression", py::overload_cast<StringOffsetExpression*>(&ModifyingVisitor::visit))
 		.def("visit_StringLengthExpression", py::overload_cast<StringLengthExpression*>(&ModifyingVisitor::visit))
 		.def("visit_NotExpression", py::overload_cast<NotExpression*>(&ModifyingVisitor::visit))
+		.def("visit_DefinedExpression", py::overload_cast<DefinedExpression*>(&ModifyingVisitor::visit))
 		.def("visit_UnaryMinusExpression", py::overload_cast<UnaryMinusExpression*>(&ModifyingVisitor::visit))
 		.def("visit_BitwiseNotExpression", py::overload_cast<BitwiseNotExpression*>(&ModifyingVisitor::visit))
 		.def("visit_AndExpression", py::overload_cast<AndExpression*>(&ModifyingVisitor::visit))
@@ -192,6 +195,7 @@ void addVisitorClasses(py::module& module)
 		.def("default_handler", static_cast<VisitResult(ModifyingVisitor::*)(const TokenStreamContext&, StringOffsetExpression*, const VisitResult&)>(&ModifyingVisitor::defaultHandler))
 		.def("default_handler", static_cast<VisitResult(ModifyingVisitor::*)(const TokenStreamContext&, StringLengthExpression*, const VisitResult&)>(&ModifyingVisitor::defaultHandler))
 		.def("default_handler", static_cast<VisitResult(ModifyingVisitor::*)(const TokenStreamContext&, NotExpression*, const VisitResult&)>(&ModifyingVisitor::defaultHandler))
+		.def("default_handler", static_cast<VisitResult(ModifyingVisitor::*)(const TokenStreamContext&, DefinedExpression*, const VisitResult&)>(&ModifyingVisitor::defaultHandler))
 		.def("default_handler", static_cast<VisitResult(ModifyingVisitor::*)(const TokenStreamContext&, UnaryMinusExpression*, const VisitResult&)>(&ModifyingVisitor::defaultHandler))
 		.def("default_handler", static_cast<VisitResult(ModifyingVisitor::*)(const TokenStreamContext&, BitwiseNotExpression*, const VisitResult&)>(&ModifyingVisitor::defaultHandler))
 		.def("default_handler", static_cast<VisitResult(ModifyingVisitor::*)(const TokenStreamContext&, AndExpression*, const VisitResult&, const VisitResult&)>(&ModifyingVisitor::defaultHandler))

diff --git a/src/python/py_visitor.h b/src/python/py_visitor.h
@@ -36,6 +36,7 @@ class PyVisitor : public yaramod::Visitor
 	PURE_VISIT(StringOffsetExpression)
 	PURE_VISIT(StringLengthExpression)
 	PURE_VISIT(NotExpression)
+	PURE_VISIT(DefinedExpression)
 	PURE_VISIT(UnaryMinusExpression)
 	PURE_VISIT(BitwiseNotExpression)
 	PURE_VISIT(AndExpression)
@@ -107,6 +108,7 @@ class PyObservingVisitor : public yaramod::ObservingVisitor
 	VISIT(ObservingVisitor, StringOffsetExpression)
 	VISIT(ObservingVisitor, StringLengthExpression)
 	VISIT(ObservingVisitor, NotExpression)
+	VISIT(ObservingVisitor, DefinedExpression)
 	VISIT(ObservingVisitor, UnaryMinusExpression)
 	VISIT(ObservingVisitor, BitwiseNotExpression)
 	VISIT(ObservingVisitor, AndExpression)
@@ -167,6 +169,7 @@ class PyModifyingVisitor : public yaramod::ModifyingVisitor
 	VISIT(ModifyingVisitor, StringOffsetExpression)
 	VISIT(ModifyingVisitor, StringLengthExpression)
 	VISIT(ModifyingVisitor, NotExpression)
+	VISIT(ModifyingVisitor, DefinedExpression)
 	VISIT(ModifyingVisitor, UnaryMinusExpression)
 	VISIT(ModifyingVisitor, BitwiseNotExpression)
 	VISIT(ModifyingVisitor, AndExpression)

diff --git a/src/python/yaramod_python.cpp b/src/python/yaramod_python.cpp
@@ -180,6 +180,7 @@ void addEnums(py::module& module)
 		.value("ImportModule", TokenType::IMPORT_MODULE)
 		.value("ImportKeyword", TokenType::IMPORT_KEYWORD)
 		.value("Not", TokenType::NOT)
+		.value("Defined", TokenType::DEFINED)
 		.value("And", TokenType::AND)
 		.value("Or", TokenType::OR)
 		.value("All", TokenType::ALL)
@@ -582,6 +583,7 @@ void addExpressionClasses(py::module& module)
 				&UnaryOpExpression::getOperand,
 				py::overload_cast<const Expression::Ptr&>(&UnaryOpExpression::setOperand));
 	unaryOpClass<NotExpression>(module, "NotExpression");
+	unaryOpClass<DefinedExpression>(module, "DefinedExpression");
 	unaryOpClass<UnaryMinusExpression>(module, "UnaryMinusExpression");
 	unaryOpClass<BitwiseNotExpression>(module, "BitwiseNotExpression");
 
@@ -816,6 +818,7 @@ void addBuilderClasses(py::module& module)
 		.def("comment", &YaraExpressionBuilder::comment, py::arg("message"), py::arg("multiline") = false, py::arg("indent") = "")
 		.def("contains", &YaraExpressionBuilder::contains)
 		.def("matches", &YaraExpressionBuilder::matches)
+		.def("defined", &YaraExpressionBuilder::defined)
 		.def("read_int8", &YaraExpressionBuilder::readInt8)
 		.def("read_int16", &YaraExpressionBuilder::readInt16)
 		.def("read_int32", &YaraExpressionBuilder::readInt32)

diff --git a/tests/cpp/builder_tests.cpp b/tests/cpp/builder_tests.cpp
@@ -1898,5 +1898,34 @@ ConjunctionWithSingleTerm) {
 )", yaraFile->getTextFormatted());
 }
 
+TEST_F(BuilderTests,
+DefinedTerm) {
+auto cond = boolVal(false).defined().get();
+
+	YaraRuleBuilder newRule;
+	auto rule = newRule
+			.withName("defined_rule")
+			.withCondition(cond)
+			.get();
+
+	YaraFileBuilder newFile;
+	auto yaraFile = newFile
+			.withRule(std::move(rule))
+			.get(true);
+
+	ASSERT_NE(nullptr, yaraFile);
+	EXPECT_EQ(R"(rule defined_rule {
+	condition:
+		defined false
+})", yaraFile->getText());
+
+	EXPECT_EQ(R"(rule defined_rule
+{
+	condition:
+		defined false
+}
+)", yaraFile->getTextFormatted());
+}
+
 }
 }
diff --git a/tests/cpp/parser_tests.cpp b/tests/cpp/parser_tests.cpp
@@ -7110,5 +7110,27 @@ rule empty_rule
 	EXPECT_EQ(input_text, driver.getParsedFile().getTextFormatted());
 }
 
+TEST_F(ParserTests,
+DefinedExpresion) {
+	prepareInput(
+R"(
+rule defined_expr
+{
+	condition:
+		defined 1
+}
+)");
+
+	EXPECT_TRUE(driver.parse(input));
+	ASSERT_EQ(1u, driver.getParsedFile().getRules().size());
+	const auto& rule = driver.getParsedFile().getRules()[0];
+
+	EXPECT_EQ("defined 1", rule->getCondition()->getText());
+	EXPECT_EQ("\"defined\"", rule->getCondition()->getFirstTokenIt()->getText());
+	EXPECT_EQ("1", rule->getCondition()->getLastTokenIt()->getPureText());
+
+	EXPECT_EQ(input_text, driver.getParsedFile().getTextFormatted());
+}
+
 }
 }
diff --git a/tests/python/test_builder.py b/tests/python/test_builder.py
@@ -1319,3 +1319,20 @@ def test_rule_with_custom_modules(self):
 	condition:
 		module_test.structure_test.function_test(/abc/) and cuckoo.sync.mutex(/abc/)
 }''')
+
+    def test_rule_with_defined_condition(self):
+        cond = yaramod.int_val(200).defined().get()
+        rule = self.new_rule \
+            .with_name('rule_with_defined_condition') \
+            .with_condition(cond.get()) \
+            .get()
+        yara_file = self.new_file \
+            .with_rule(rule) \
+            .get(True)
+
+        self.assertEqual(yara_file.text_formatted, '''rule rule_with_defined_condition
+{
+	condition:
+		defined 200
+}
+''')