Merge pull request #251 from vojone/string_literal_sets

String literal sets

src/parser/parser_driver.cpp

@@ -1141,7 +1141,7 @@ void ParserDriver::defineGrammar()
 			return output;
 		})
 		.production(
-			"FOR", "for_expression", "ID", "IN", "integer_set", [&](auto&& args) -> Value {
+			"FOR", "for_expression", "ID", "IN", "for_expression_set", [&](auto&& args) -> Value {
 				auto symbol = std::make_shared<ValueSymbol>(args[2].getTokenIt()->getString(), Expression::Type::Int);
 				if (!addLocalSymbol(symbol))
 					error_handle(args[2].getTokenIt()->getLocation(), "Redefinition of identifier '" + args[2].getTokenIt()->getString() + "'");
@@ -2143,8 +2143,8 @@ void ParserDriver::defineGrammar()
 		})
 		;
 
-	_parser.rule("integer_set") // Expression::Ptr
-		.production("LP", "integer_enumeration", "RP", [&](auto&& args) -> Value {
+	_parser.rule("for_expression_set") // Expression::Ptr
+		.production("LP", "for_expression_enumeration", "RP", [&](auto&& args) -> Value {
 			auto lp = args[0].getTokenIt();
 			auto rp = args[2].getTokenIt();
 			lp->setType(TokenType::LP_ENUMERATION);
@@ -2159,18 +2159,20 @@ void ParserDriver::defineGrammar()
 		})
 		;
 
-	_parser.rule("integer_enumeration") // vector<Expression::Ptr>
+	_parser.rule("for_expression_enumeration") // vector<Expression::Ptr>
 		.production("primary_expression", [&](auto&& args) -> Value {
 			auto expr = args[0].getExpression();
-			if (!expr->isInt())
-				error_handle(currentFileContext()->getLocation(), "integer set expects integer type");
+			if (!expr->isInt() && !expr->isString())
+				error_handle(currentFileContext()->getLocation(), "for expression set expects string or integer type");
 			return std::vector<Expression::Ptr>{std::move(expr)};
 		})
-		.production("integer_enumeration", "COMMA", "primary_expression", [&](auto&& args) -> Value {
-			auto expr = args[2].getExpression();
-			if (!expr->isInt())
-				error_handle(currentFileContext()->getLocation(), "integer set expects integer type");
+		.production("for_expression_enumeration", "COMMA", "primary_expression", [&](auto&& args) -> Value {
 			auto output = std::move(args[0].getMultipleExpressions());
+			auto expr = args[2].getExpression();
+			if (!expr->isInt() && !expr->isString())
+				error_handle(currentFileContext()->getLocation(), "for expression set expects string or integer type");
+			if(output[0]->getType() != expr->getType())
+				error_handle(currentFileContext()->getLocation(), "all items in enumeration must be same type");
 			output.push_back(std::move(expr));
 			return output;
 		})

tests/cpp/parser_tests.cpp

@@ -3097,6 +3097,104 @@ rule for_string_set_condition
 	EXPECT_EQ(input_text, driver.getParsedFile().getTextFormatted());
 }
 
+TEST_F(ParserTests,
+ForStringLiteralSetConditionWorks) {
+	prepareInput(
+R"(
+import "pe"
+
+rule for_string_literal_set_condition
+{
+	condition:
+		for any s in ("hash1", "hash2", "hash3") : ( pe.imphash() == s )
+}
+)");
+
+	EXPECT_TRUE(driver.parse(input));
+	ASSERT_EQ(1u, driver.getParsedFile().getRules().size());
+
+	const auto& rule = driver.getParsedFile().getRules()[0];
+	EXPECT_EQ("for any s in (\"hash1\", \"hash2\", \"hash3\") : ( pe.imphash() == s )", rule->getCondition()->getText());
+	EXPECT_EQ("for", rule->getCondition()->getFirstTokenIt()->getPureText());
+	EXPECT_EQ(")", rule->getCondition()->getLastTokenIt()->getPureText());
+
+	EXPECT_EQ(input_text, driver.getParsedFile().getTextFormatted());
+}
+
+TEST_F(ParserTests,
+ForStringLiteralSetWithOneStringConditionWorks) {
+	prepareInput(
+R"(
+import "pe"
+
+rule for_string_literal_set_condition
+{
+	condition:
+		for any s in ("hash1") : ( pe.imphash() == s )
+}
+)");
+
+	EXPECT_TRUE(driver.parse(input));
+	ASSERT_EQ(1u, driver.getParsedFile().getRules().size());
+
+	const auto& rule = driver.getParsedFile().getRules()[0];
+	EXPECT_EQ("for any s in (\"hash1\") : ( pe.imphash() == s )", rule->getCondition()->getText());
+	EXPECT_EQ("for", rule->getCondition()->getFirstTokenIt()->getPureText());
+	EXPECT_EQ(")", rule->getCondition()->getLastTokenIt()->getPureText());
+
+	EXPECT_EQ(input_text, driver.getParsedFile().getTextFormatted());
+}
+
+TEST_F(ParserTests,
+ForStringLiteralSetWithStringSymbolsConditionWorks) {
+	prepareInput(
+R"(
+import "pe"
+import "elf"
+
+rule for_string_literal_set_condition
+{
+	condition:
+		for any s in ("hash1", elf.dynsym [ 0 ].name, pe.imphash ( )) : ( "abc123" == s )
+}
+)");
+
+	EXPECT_TRUE(driver.parse(input));
+	ASSERT_EQ(1u, driver.getParsedFile().getRules().size());
+
+	const auto& rule = driver.getParsedFile().getRules()[0];
+	EXPECT_EQ("for any s in (\"hash1\", elf.dynsym[0].name, pe.imphash()) : ( \"abc123\" == s )", rule->getCondition()->getText());
+	EXPECT_EQ("for", rule->getCondition()->getFirstTokenIt()->getPureText());
+	EXPECT_EQ(")", rule->getCondition()->getLastTokenIt()->getPureText());
+
+	EXPECT_EQ(input_text, driver.getParsedFile().getTextFormatted());
+}
+
+TEST_F(ParserTests,
+ForExpressionSetWithItemsVariousTypesForbidden) {
+	prepareInput(
+R"(
+import "pe"
+
+rule for_string_literal_set_condition
+{
+	condition:
+		for any s in ("hash1", 1, "hash3") : ( pe.imphash() == s )
+}
+)");
+
+	try
+	{
+		driver.parse(input);
+		FAIL() << "Parser did not throw an exception.";
+	}
+	catch (const ParserError& err)
+	{
+		EXPECT_EQ(0u, driver.getParsedFile().getRules().size());
+		EXPECT_EQ("Error at 7.27: all items in enumeration must be same type", err.getErrorMessage());
+	}
+}
+
 TEST_F(ParserTests,
 NoneOfThemConditionWorks) {
 	prepareInput(

tests/python/test_builder.py

@@ -1367,6 +1367,37 @@ def test_rule_with_complex_condition(self):
 	condition:
 		for any i in (1, 2, 3) : ( $1 at (entrypoint + i) )
 }''')
+
+    def test_rule_with_string_literal_set(self):
+        cond = yaramod.for_loop(
+                yaramod.any(),
+                's',
+                yaramod.set([
+                    yaramod.string_val('hash1'),
+                    yaramod.string_val('hash2'),
+                    yaramod.string_val('hash3')
+                ]),
+                yaramod.id('s') == yaramod.string_val('abc123')
+            )
+        rule = self.new_rule \
+            .with_name('rule_with_string_literal_set') \
+            .with_condition(cond.get()) \
+            .get()
+        yara_file = self.new_file \
+            .with_rule(rule) \
+            .get()
+
+
+        self.assertEqual(yara_file.text_formatted, '''rule rule_with_string_literal_set
+{
+	condition:
+		for any s in ("hash1", "hash2", "hash3") : ( s == "abc123" )
+}
+''')
+        self.assertEqual(yara_file.text, '''rule rule_with_string_literal_set {
+	condition:
+		for any s in ("hash1", "hash2", "hash3") : ( s == "abc123" )
+}''')
 
     def test_rule_with_string_modifiers(self):
         rule = self.new_rule \

tests/python/test_parser.py

@@ -1173,6 +1173,24 @@ def test_for_integer_set_condition(self):
         self.assertTrue(isinstance(rule.condition.body, yaramod.IdExpression))
         self.assertEqual(rule.condition.text, 'for all i in (1, 2, 3) : ( i )')
 
+    def test_for_string_literal_set_condition(self):
+        yara_file = yaramod.Yaramod().parse_string('''
+import "pe"
+
+rule for_integer_set_condition {
+    condition:
+        for all i in ("hash1","hash2","hash3") : ( i == pe.imphash() )
+}''')
+
+        self.assertEqual(len(yara_file.rules), 1)
+
+        rule = yara_file.rules[0]
+        self.assertTrue(isinstance(rule.condition, yaramod.ForArrayExpression))
+        self.assertTrue(isinstance(rule.condition.variable, yaramod.AllExpression))
+        self.assertTrue(isinstance(rule.condition.iterable, yaramod.SetExpression))
+        self.assertTrue(isinstance(rule.condition.body, yaramod.EqExpression))
+        self.assertEqual(rule.condition.text, 'for all i in ("hash1", "hash2", "hash3") : ( i == pe.imphash() )')
+
     def test_for_array_condition(self):
         yara_file = yaramod.Yaramod().parse_string('''
 import "pe"