Repository: avast/yaramod (language: python)

Issue:
Make commenting of builded Expressions more flexible.
We need a way to comment created Expressions from behind and an option to do it on the same line with a one-line comment. Also when creating conjunctions or disjunctions we would appreciate if it was possible for the comments to be on the same line as the particular terms of the operation and not to place the comments on separate lines. This issue is being implemented in #194.


Discussion / hints:

Fix patch:
diff --git a/docs/rtd/creating_rulesets.rst b/docs/rtd/creating_rulesets.rst
index 7cfe500b..e61f2ab4 100644
--- a/docs/rtd/creating_rulesets.rst
+++ b/docs/rtd/creating_rulesets.rst
@@ -337,8 +337,8 @@ basic expressions and find the most suitable one.
         * ``forLoop(spec, set, body)`` - represents ``for`` loop over set of string references (``forLoop(any(), set({stringRef("$*")}), matchAt("$", intVal(100))``)
         * ``of(spec, set)`` - represents ``<spec> of <set>`` (``of(all(), them())``)
         * ``paren(expr, [newline])`` - represents parentheses around expressions and ``newline`` indicator for putting enclosed expression on its own line (``paren(intVal(10))``)
-        * ``conjunction(terms, [newline])`` - represents conjunction of ``terms`` and optionally puts them on each separate line if ``newline`` is set (``conjunction({id("rule1"), id("rule2")})``)
-        * ``disjunction(terms, [newline])`` - represents disjunction of ``terms`` and optionally puts them on each separate line if ``newline`` is set (``disjunction({id("rule1"), id("rule2")})``)
+        * ``conjunction(terms, [newline])`` - represents conjunction of ``terms`` and optionally puts them on each separate line if ``newline`` is set (``conjunction({id("rule1"), id("rule2")})``). The ``terms`` parameter can be an array containing other expressions to be put together in the conjunction. But also ``terms`` can be an array of pairs, where each pair contains a term to be put in the conjunction and a comment, which will be associated with the term and printed on the same line
+        * ``disjunction(terms, [newline])`` - represents disjunction of ``terms`` and optionally puts them on each separate line if ``newline`` is set (``disjunction({id("rule1"), id("rule2")})``). The ``terms`` parameter can be an array containing other expressions to be put together in the disjunction. But also ``terms`` can be an array of pairs, where each pair contains a term to be put in the disjunction and a comment, which will be associated with the term and printed on the same line
 
         **Complex expression methods**
 
@@ -374,6 +374,8 @@ basic expressions and find the most suitable one.
         * ``readUInt8(be)`` - represents call to special function ``uint8(be)`` (``intVal(100).readUInt8()``)
         * ``readUInt16(be)`` - represents call to special function ``uint16(be)`` (``intVal(100).readUInt16()``)
         * ``readUInt32(be)`` - represents call to special function ``uint32(be)`` (``intVal(100).readUInt32()``)
+        * ``comment(message, [multiline], [indent], [linebreak])`` - adds a comment ``message`` to the expression which then appears in the formatted text before the expression. Only the ``message`` parameter is required, the ``multiline`` (default ``false``), ``indent`` (default "") and ``linebreak`` (default ``true``) parameters are optional
+        * ``commentBehind(message, [multiline], [indent], [linebreak])`` - adds a comment ``message`` to the expression which then appears in the formatted text after the expression. Only the ``message`` parameter is required, the ``multiline`` (default ``false``), ``indent`` (default "") and ``linebreak`` (default ``true``) parameters are optional
 
         **Hex strings**
 
diff --git a/include/yaramod/builder/yara_expression_builder.h b/include/yaramod/builder/yara_expression_builder.h
index e52cc79c..91e35472 100644
--- a/include/yaramod/builder/yara_expression_builder.h
+++ b/include/yaramod/builder/yara_expression_builder.h
@@ -8,6 +8,7 @@
 
 #include <memory>
 #include <string>
+#include <iostream>
 
 #include "yaramod/types/expression.h"
 #include "yaramod/types/token_stream.h"
@@ -166,7 +167,8 @@ class YaraExpressionBuilder
 	YaraExpressionBuilder& operator<<(const YaraExpressionBuilder& other);
 	YaraExpressionBuilder& operator>>(const YaraExpressionBuilder& other);
 
-	YaraExpressionBuilder& comment(const std::string& message, bool multiline = false, const std::string& indent = {});
+	YaraExpressionBuilder& comment(const std::string& message, bool multiline = false, const std::string& indent = {}, bool linebreak = true);
+	YaraExpressionBuilder& commentBehind(const std::string& message, bool multiline = false, const std::string& indent = {}, bool linebreak = true);
 	YaraExpressionBuilder& call(const std::vector<YaraExpressionBuilder>& args);
 	/**
 	 * Calls function from an expression
@@ -255,6 +257,8 @@ YaraExpressionBuilder range(const YaraExpressionBuilder& low, const YaraExpressi
 
 YaraExpressionBuilder conjunction(const YaraExpressionBuilder& lhs, const YaraExpressionBuilder& rhs, bool linebreak = false);
 YaraExpressionBuilder disjunction(const YaraExpressionBuilder& lhs, const YaraExpressionBuilder& rhs, bool linebreak = false);
+YaraExpressionBuilder conjunction(const YaraExpressionBuilder& lhs, const std::string& lhscomment, const YaraExpressionBuilder& rhs);
+YaraExpressionBuilder disjunction(const YaraExpressionBuilder& lhs, const std::string& lhscomment, const YaraExpressionBuilder& rhs);
 YaraExpressionBuilder conjunction(const std::vector<YaraExpressionBuilder>& terms, bool linebreaks = false);
 YaraExpressionBuilder disjunction(const std::vector<YaraExpressionBuilder>& terms, bool linebreaks = false);
 YaraExpressionBuilder conjunction(const std::vector<std::pair<YaraExpressionBuilder, std::string>>& terms);
diff --git a/include/yaramod/types/token_stream.h b/include/yaramod/types/token_stream.h
index 7efc7299..8f6bf59f 100644
--- a/include/yaramod/types/token_stream.h
+++ b/include/yaramod/types/token_stream.h
@@ -87,6 +87,14 @@ class TokenStream
 	TokenIt erase(TokenIt element);
 	TokenIt erase(TokenIt first, TokenIt last);
 
+
+	// Puts comment to the front.
+	void comment(const std::string& message, bool multiline = false, const std::string& indent = {}, bool linebreak = true);
+	// Puts comment to the back.
+	void commentBehind(const std::string& message, bool multiline = false, const std::string& indent = {}, bool linebreak = true);
+	// Puts comment before the insert_before parameter token.
+	void commentBeforeToken(const std::string& message, TokenIt insert_before, bool multiline = false, const std::string& indent = {}, bool linebreak = true);
+
 	// Steals all data from donor and append it at the end.
 	void moveAppend(TokenStream* donor);
 	// Steals all data from donor and append it at position before.
diff --git a/src/builder/yara_expression_builder.cpp b/src/builder/yara_expression_builder.cpp
index 696908a5..a68293d4 100644
--- a/src/builder/yara_expression_builder.cpp
+++ b/src/builder/yara_expression_builder.cpp
@@ -55,6 +55,31 @@ YaraExpressionBuilder logicalFormula(std::vector<YaraExpressionBuilder> terms, c
 	return formula;
 }
 
+template <typename Op>
+YaraExpressionBuilder logicalFormula(std::vector<YaraExpressionBuilder> terms, std::vector<std::string> comments, const Op& op)
+{
+	if (terms.empty())
+		return boolVal(true);
+
+	if (terms.size() == 1)
+		return terms.front();
+
+	auto formula = op(terms[0], comments[0], terms[1]);
+	for (std::size_t i = 2; i < terms.size(); ++i)
+	{
+		if (!terms[i].canBeBool())
+		{
+			const auto& expr = terms[i].get();
+			error_handle("Expected boolean, got '" + expr->getText() + "' of type " + expr->getTypeString());
+		}
+		if (i >= 2)
+			formula = op(formula, comments[i-1], terms[i]);
+	}
+
+	formula.setType(Expression::Type::Bool);
+	return formula;
+}
+
 } //namespace
 
 /**
@@ -459,39 +484,24 @@ YaraExpressionBuilder& YaraExpressionBuilder::call(const std::vector<YaraExpress
  *
  * @return Builder.
  */
-YaraExpressionBuilder& YaraExpressionBuilder::comment(const std::string& message, bool multiline, const std::string& indent)
+YaraExpressionBuilder& YaraExpressionBuilder::comment(const std::string& message, bool multiline, const std::string& indent, bool linebreak)
 {
-	if (!message.empty())
-	{
-		TokenIt insert_before = _tokenStream->begin();
-		std::stringstream ss;
-		ss << indent;
-		if (multiline)
-		{
-			ss << "/*";
-			if (message.front() != '\n')
-				ss << " ";
-			for (auto c : message)
-			{
-				ss << c;
-				if (c == '\n')
-					ss << indent;
-			}
-			if (message.back() != '\n')
-				ss << " ";
-			ss << "*/";
-			_tokenStream->emplace(insert_before, TokenType::COMMENT, ss.str());
-		}
-		else
-		{
-			for (auto item : message)
-				if (item == '\n')
-					throw YaraExpressionBuilderError("Error: one-line comment must not contain \\n.");
-			ss << "// " << message;
-			_tokenStream->emplace(insert_before, TokenType::ONELINE_COMMENT, ss.str());
-		}
-		_tokenStream->emplace(insert_before, TokenType::NEW_LINE, "\n");
-	}
+	_tokenStream->comment(message, multiline, indent, linebreak);
+	return *this;
+}
+
+/**
+ * Puts comment behind the expression.
+ *
+ * @param message The comment message.
+ * @param multiline If set, the commet will be multiline.
+ * @param indent Additional indent added to the indentation computed by the autoformatter.
+ *
+ * @return Builder.
+ */
+YaraExpressionBuilder& YaraExpressionBuilder::commentBehind(const std::string& message, bool multiline, const std::string& indent, bool linebreak)
+{
+	_tokenStream->commentBehind(message, multiline, indent, linebreak);
 	return *this;
 }
 
@@ -1205,6 +1215,58 @@ YaraExpressionBuilder disjunction(const YaraExpressionBuilder& lhs, const YaraEx
 	return YaraExpressionBuilder(std::move(ts), std::move(expression), Expression::Type::Bool);
 }
 
+/**
+ * Creates conjunction.
+ *
+ * @param lhs Left-hand side.
+ * @param lhscomment First comment.
+ * @param rhs Right-hand side.
+ *
+ * @return Builder.
+ */
+YaraExpressionBuilder conjunction(const YaraExpressionBuilder& lhs, const std::string& lhscomment, const YaraExpressionBuilder& rhs)
+{
+	if (!lhs.canBeBool())
+		error_handle(ArgType::Left, "and", "bool", rhs.get());
+	else if (!rhs.canBeBool())
+		error_handle(ArgType::Right, "and", "bool", lhs.get());
+
+	auto ts = std::make_shared<TokenStream>();
+	ts->moveAppend(lhs.getTokenStream());
+	TokenIt andToken = ts->emplace_back(TokenType::AND, "and");
+	ts->commentBehind(lhscomment, false, "", true);
+	ts->moveAppend(rhs.getTokenStream());
+
+	auto expression = std::make_shared<AndExpression>(lhs.get(), andToken, rhs.get(), true);
+	return YaraExpressionBuilder(std::move(ts), std::move(expression), Expression::Type::Bool);
+}
+
+/**
+ * Creates disjunction.
+ *
+ * @param lhs Left-hand side.
+ * @param lhscomment First comment.
+ * @param rhs Right-hand side.
+ *
+ * @return Builder.
+ */
+YaraExpressionBuilder disjunction(const YaraExpressionBuilder& lhs, const std::string& lhscomment, const YaraExpressionBuilder& rhs)
+{
+	if (!lhs.canBeBool())
+		error_handle(ArgType::Left, "or", "bool", lhs.get());
+	else if (!rhs.canBeBool())
+		error_handle(ArgType::Right, "or", "bool", rhs.get());
+
+	auto ts = std::make_shared<TokenStream>();
+	ts->moveAppend(lhs.getTokenStream());
+	TokenIt orToken = ts->emplace_back(TokenType::OR, "or");
+	ts->commentBehind(lhscomment, false, "", true);
+	ts->moveAppend(rhs.getTokenStream());
+
+	auto expression = std::make_shared<OrExpression>(lhs.get(), orToken, rhs.get(), true);
+	return YaraExpressionBuilder(std::move(ts), std::move(expression), Expression::Type::Bool);
+}
+
 /**
  * Creates conjunction of terms.
  *
@@ -1245,17 +1307,23 @@ YaraExpressionBuilder disjunction(const std::vector<YaraExpressionBuilder>& term
  *
  * @return Builder.
  */
-YaraExpressionBuilder conjunction(const std::vector<std::pair<YaraExpressionBuilder, std::string>>& terms)
+YaraExpressionBuilder conjunction(const std::vector<std::pair<YaraExpressionBuilder, std::string>>& commented_terms)
 {
-	std::vector<YaraExpressionBuilder> commented_terms;
-	for (const auto& pair : terms)
+	std::vector<YaraExpressionBuilder> terms;
+	std::vector<std::string> comments;
+	for (const auto& pair : commented_terms)
 	{
 		if (!pair.first.canBeBool())
 			error_handle(ArgType::Single, "and", "bool", pair.first.get());
-		commented_terms.push_back(pair.first);
-		commented_terms.back().comment(pair.second, true);
+		terms.push_back(pair.first);
+		comments.push_back(pair.second);
 	}
-	return logicalFormula(commented_terms, [](YaraExpressionBuilder& term1, YaraExpressionBuilder& term2) { return conjunction(term1, term2, true); });
+	auto output = logicalFormula(terms, comments, []( YaraExpressionBuilder& term1, std::string& comment1, YaraExpressionBuilder& term2) {
+		return conjunction(term1, comment1, term2);
+	});
+
+	output.commentBehind(comments.back(), false, "", false);
+	return output;
 }
 
 /**
@@ -1266,17 +1334,23 @@ YaraExpressionBuilder conjunction(const std::vector<std::pair<YaraExpressionBuil
  *
  * @return Builder.
  */
-YaraExpressionBuilder disjunction(const std::vector<std::pair<YaraExpressionBuilder, std::string>>& terms)
+YaraExpressionBuilder disjunction(const std::vector<std::pair<YaraExpressionBuilder, std::string>>& commented_terms)
 {
-	std::vector<YaraExpressionBuilder> commented_terms;
-	for (const auto& pair : terms)
+	std::vector<YaraExpressionBuilder> terms;
+	std::vector<std::string> comments;
+	for (const auto& pair : commented_terms)
 	{
 		if (!pair.first.canBeBool())
-			error_handle(ArgType::Single, "or", "bool", pair.first.get());
-		commented_terms.push_back(pair.first);
-		commented_terms.back().comment(pair.second, true);
+			error_handle(ArgType::Single, "and", "bool", pair.first.get());
+		terms.push_back(pair.first);
+		comments.push_back(pair.second);
 	}
-	return logicalFormula(commented_terms, [](YaraExpressionBuilder& term1, YaraExpressionBuilder& term2) { return disjunction(term1, term2, true); });
+	auto output = logicalFormula(terms, comments, []( YaraExpressionBuilder& term1, std::string& comment1, YaraExpressionBuilder& term2) {
+		return disjunction(term1, comment1, term2);
+	});
+
+	output.commentBehind(comments.back(), false, "", false);
+	return output;
 }
 
 /**
diff --git a/src/python/yaramod_python.cpp b/src/python/yaramod_python.cpp
index 8e4d1551..8f023a8b 100644
--- a/src/python/yaramod_python.cpp
+++ b/src/python/yaramod_python.cpp
@@ -526,7 +526,8 @@ void addTokenStreamClass(py::module& module)
 		.def_property_readonly("front", &TokenStream::front)
 		.def_property_readonly("back", &TokenStream::back)
 		.def_property_readonly("tokens", &TokenStream::getTokens)
-		.def_property_readonly("tokens_as_text", &TokenStream::getTokensAsText);
+		.def_property_readonly("tokens_as_text", &TokenStream::getTokensAsText)
+		.def("comment_before_token", &TokenStream::commentBeforeToken, py::arg("message"), py::arg("insert_before"), py::arg("multiline") = false, py::arg("indent") = "", py::arg("linebreak") = true);
 }
 
 void addExpressionClasses(py::module& module)
@@ -821,7 +822,8 @@ void addBuilderClasses(py::module& module)
 			})
 		.def("__getitem__", &YaraExpressionBuilder::operator[])
 		.def("access", &YaraExpressionBuilder::access)
-		.def("comment", &YaraExpressionBuilder::comment, py::arg("message"), py::arg("multiline") = false, py::arg("indent") = "")
+		.def("comment", &YaraExpressionBuilder::comment, py::arg("message"), py::arg("multiline") = false, py::arg("indent") = "", py::arg("linebreak") = true)
+		.def("comment_behind", &YaraExpressionBuilder::commentBehind, py::arg("message"), py::arg("multiline") = false, py::arg("indent") = "", py::arg("linebreak") = true)
 		.def("contains", &YaraExpressionBuilder::contains)
 		.def("matches", &YaraExpressionBuilder::matches)
 		.def("iequals", &YaraExpressionBuilder::iequals)
diff --git a/src/types/token_stream.cpp b/src/types/token_stream.cpp
index 3a0ae749..ade5667f 100644
--- a/src/types/token_stream.cpp
+++ b/src/types/token_stream.cpp
@@ -490,6 +490,62 @@ void TokenStream::removeRedundantDoubleNewlines()
 	}
 }
 
+void TokenStream::comment(const std::string& message, bool multiline, const std::string& indent, bool linebreak)
+{
+	auto insert_before = empty() ? end() : begin();
+	commentBeforeToken(message, insert_before, multiline, indent, linebreak);
+}
+
+void TokenStream::commentBehind(const std::string& message, bool multiline, const std::string& indent, bool linebreak)
+{
+	auto insert_before = end();
+	while (insert_before != begin())
+	{
+		auto predecessor = std::prev(insert_before);
+		if (predecessor->getType() == TokenType::NEW_LINE)
+			--insert_before;
+		else
+			break;
+	}
+	commentBeforeToken(message, insert_before, multiline, indent, linebreak);
+}
+
+void TokenStream::commentBeforeToken(const std::string& message, TokenIt insert_before, bool multiline, const std::string& indent, bool linebreak)
+{
+	if (!message.empty())
+	{
+		std::stringstream ss;
+		ss << indent;
+		if (multiline)
+		{
+			ss << "/*";
+			if (message.front() != '\n')
+				ss << " ";
+			for (auto c : message)
+			{
+				ss << c;
+				if (c == '\n')
+					ss << indent;
+			}
+			if (message.back() != '\n')
+				ss << " ";
+			ss << "*/";
+			emplace(insert_before, TokenType::COMMENT, ss.str());
+		}
+		else
+		{
+			for (auto item : message)
+				if (item == '\n')
+					throw YaramodError("Error: one-line comment must not contain \\n.");
+			ss << "// " << message;
+			emplace(insert_before, TokenType::ONELINE_COMMENT, ss.str());
+		}
+		// if (insert_before != end() && (linebreak || !multiline))
+		if (linebreak)
+			emplace(insert_before, TokenType::NEW_LINE, "\n");
+	}
+}
+
 void TokenStream::addMissingNewLines()
 {
 	BracketStack brackets;


Test patch:
diff --git a/tests/cpp/builder_tests.cpp b/tests/cpp/builder_tests.cpp
index b7fdbd3e..d1f256b6 100644
--- a/tests/cpp/builder_tests.cpp
+++ b/tests/cpp/builder_tests.cpp
@@ -879,7 +879,7 @@ rule rule_with_range
 
 TEST_F(BuilderTests,
 RuleWithConditionWithOnelineComment) {
-	auto cond = (id("pe").comment("Number of sections needs to exceed 1", false).access("number_of_sections") > intVal(1)).get();
+	auto cond = (id("pe").access("number_of_sections") > intVal(1)).comment("Number of sections needs to exceed 1", false).get();
 	EXPECT_EQ("pe", cond->getFirstTokenIt()->getPureText());
 	EXPECT_EQ("1", cond->getLastTokenIt()->getPureText());
 
@@ -916,7 +916,7 @@ rule rule_with_range
 
 TEST_F(BuilderTests,
 RuleWithConditionWithMultilineComment) {
-	auto cond = (id("pe").comment("Number of sections needs to exceed 1,\n\t\tbecause one is simply not enough.", true).access("number_of_sections") > intVal(1)).get();
+	auto cond = (id("pe").access("number_of_sections") > intVal(1)).comment("Number of sections needs to exceed 1,\n\t\tbecause one is simply not enough.", true).get();
 
 	YaraRuleBuilder newRule;
 	auto rule = newRule
@@ -1537,12 +1537,9 @@ RuleWithCommentedConjunctionInConditionWorks) {
 	strings:
 		$1 = "Hello"
 	condition:
-		/* comment1 */
-		$1 and
-		/* comment2 */
-		(@1 < 100) and
-		/* comment3 */
-		(entrypoint == 100)
+		$1 and              // comment1
+		(@1 < 100) and      // comment2
+		(entrypoint == 100) // comment3
 }
 )", yaraFile->getTextFormatted());
 }
@@ -1588,12 +1585,9 @@ RuleWithCommentedDisjunctionInConditionWorks) {
 	condition:
 		true and
 		(
-			/* Hello must be present */
-			$1 or
-			/* World must be present */
-			$2 or
-			/* entrypoint is 100 */
-			(entrypoint == 100)
+			$1 or               // Hello must be present
+			$2 or               // World must be present
+			(entrypoint == 100) // entrypoint is 100
 		)
 }
 )", yaraFile->getTextFormatted());
diff --git a/tests/python/test_builder.py b/tests/python/test_builder.py
index 6900f851..bcf6f915 100644
--- a/tests/python/test_builder.py
+++ b/tests/python/test_builder.py
@@ -911,10 +911,31 @@ def test_rule_with_and_condition_with_comments(self):
         self.assertEqual(yara_file.text_formatted, '''rule rule_with_and_condition_with_comments
 {
 	condition:
-		/* comment1 */
+		filesize > 100 and // comment1
+		filesize < 200     // comment2
+}
+''')
+        self.assertEqual(yara_file.text, '''rule rule_with_and_condition_with_comments {
+	condition:
 		filesize > 100 and
-		/* comment2 */
 		filesize < 200
+}''')
+
+    def test_rule_with_and_condition_with_comments_behind(self):
+        cond = yaramod.conjunction([[yaramod.filesize() > yaramod.int_val(100), 'comment1'], [yaramod.filesize() < yaramod.int_val(200), 'comment2']])
+        rule = self.new_rule \
+            .with_name('rule_with_and_condition_with_comments') \
+            .with_condition(cond.get()) \
+            .get()
+        yara_file = self.new_file \
+            .with_rule(rule) \
+            .get()
+
+        self.assertEqual(yara_file.text_formatted, '''rule rule_with_and_condition_with_comments
+{
+	condition:
+		filesize > 100 and // comment1
+		filesize < 200     // comment2
 }
 ''')
         self.assertEqual(yara_file.text, '''rule rule_with_and_condition_with_comments {
@@ -936,10 +957,8 @@ def test_rule_with_or_condition_with_comments(self):
         self.assertEqual(yara_file.text_formatted, '''rule rule_with_or_condition_with_comments
 {
 	condition:
-		/* skip small files */
-		filesize > 100 or
-		/* also too big files */
-		filesize < 200
+		filesize > 100 or // skip small files
+		filesize < 200    // also too big files
 }
 ''')
         self.assertEqual(yara_file.text, '''rule rule_with_or_condition_with_comments {
@@ -1075,6 +1094,34 @@ def test_rule_with_function_call_condition(self):
 ''')
         self.assertEqual(yara_file.text, '''import "pe"
 
+rule rule_with_function_call_condition {
+	condition:
+		pe.is_dll()
+}''')
+
+    def test_rule_with_function_call_and_oneline_comment(self):
+        cond = yaramod.id('pe').access('is_dll')().comment_behind(message="Generated", multiline=False, linebreak=False)
+        rule = self.new_rule \
+            .with_name('rule_with_function_call_condition') \
+            .with_condition(cond.get()) \
+            .get()
+        yara_file = self.new_file \
+            .with_module('pe') \
+            .with_rule(rule) \
+            .get()
+
+        print(yara_file.text_formatted)
+
+        self.assertEqual(yara_file.text_formatted, '''import "pe"
+
+rule rule_with_function_call_condition
+{
+	condition:
+		pe.is_dll() // Generated
+}
+''')
+        self.assertEqual(yara_file.text, '''import "pe"
+
 rule rule_with_function_call_condition {
 	condition:
 		pe.is_dll()