feat: Added support for 'with' expression

include/yaramod/builder/yara_expression_builder.h

@@ -280,6 +280,9 @@ YaraExpressionBuilder none();
 YaraExpressionBuilder them();
 
 YaraExpressionBuilder regexp(const std::string& text, const std::string& suffixMods = std::string{});
+
+YaraExpressionBuilder var_def(const std::string& name, const YaraExpressionBuilder& expr);
+YaraExpressionBuilder with(const std::vector<YaraExpressionBuilder>& vars, const YaraExpressionBuilder& body);
 /// @}
 
 }

include/yaramod/types/expressions.h

@@ -2407,4 +2407,110 @@ class RegexpExpression : public Expression
 	std::shared_ptr<String> _regexp; ///< Regular expression string
 };
 
+
+/**
+ * Class representing variable definition within with expression.
+ *
+ * For example:
+ * @code
+ * with last_section = pe.sections[pe.number_of_sections - 1] : ( ... )
+ *      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+ * @endcode
+ */
+class VariableDefExpression : public Expression
+{
+public:
+	template <typename ExpPtr>
+	VariableDefExpression(TokenIt name, ExpPtr&& expr)
+		: _name(name),
+		  _expr(std::forward<ExpPtr>(expr))
+	{
+	}
+
+	virtual VisitResult accept(Visitor* v) override
+	{
+		return v->visit(this);
+	}
+
+	virtual TokenIt getFirstTokenIt() const override { return _name; }
+	virtual TokenIt getLastTokenIt() const override { return _expr->getLastTokenIt(); }
+
+	const std::string& getName() const { return _name->getString(); }
+	const Expression::Ptr& getExpression() const { return _expr; }
+
+	virtual std::string getText(const std::string& indent = std::string{}) const override
+	{
+		return getName() + " = " + _expr->getText(indent);
+	}
+
+	void setExpression(const Expression::Ptr& expr) { _expr = expr; }
+	void setExpression(Expression::Ptr&& expr) { _expr = std::move(expr); }
+
+private:
+	TokenIt _name;
+	Expression::Ptr _expr;
+};
+
+/**
+ * Class representing with variable expression.
+ *
+ * For example:
+ * @code
+ * with last_section = pe.sections[pe.number_of_sections - 1] : ( ... )
+ * ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+ * @endcode
+ */
+class WithExpression : public Expression
+{
+public:
+	template <typename VarVector, typename ExpPtr>
+	WithExpression(TokenIt with, VarVector&& vars, ExpPtr&& body, TokenIt right_bracket)
+		: _with(with),
+		  _vars(std::forward<VarVector>(vars)),
+		  _body(std::forward<ExpPtr>(body)),
+		  _right_bracket(right_bracket)
+	{
+	}
+
+	virtual VisitResult accept(Visitor* v) override
+	{
+		return v->visit(this);
+	}
+
+	virtual TokenIt getFirstTokenIt() const override { return _with; }
+	virtual TokenIt getLastTokenIt() const override { return _right_bracket; }
+
+	const std::vector<Expression::Ptr>& getVariables() const { return _vars; }
+	const Expression::Ptr& getBody() const { return _body; }
+
+	virtual std::string getText(const std::string& indent = std::string{}) const override
+	{
+		std::ostringstream ss;
+		ss << "with ";
+		for (auto itr = _vars.begin(), end = _vars.end(); itr != end; ++itr)
+		{
+			auto& expr = *itr;
+			ss << expr->getText(indent);
+			if (itr + 1 != end)
+				ss << ", ";
+			else
+				ss << " : ";
+		}
+		ss << "(" << _body->getText(indent) << ")";
+		return ss.str();
+	}
+
+	void setVariables(const std::vector<Expression::Ptr>& vars) { _vars = vars; }
+	void setVariables(std::vector<Expression::Ptr>&& vars) { _vars = std::move(vars); }
+
+	void setBody(const Expression::Ptr& body) { _body = body; }
+	void setBody(Expression::Ptr&& body) { _body = std::move(body); }
+
+private:
+	TokenIt _with;
+	std::vector<Expression::Ptr> _vars;
+	Expression::Ptr _body;
+	TokenIt _right_bracket;
+};
+
 }

include/yaramod/types/token_type.h

@@ -156,6 +156,7 @@ enum class TokenType
 	INCLUDE_PATH,
 	FUNCTION_CALL_LP,
 	FUNCTION_CALL_RP,
+	WITH,
 	INVALID,
 };
 

include/yaramod/utils/modifying_visitor.h

@@ -427,6 +427,27 @@ class ModifyingVisitor : public Visitor
 	}
 
 	virtual VisitResult visit(RegexpExpression*) override { return {}; }
+
+	virtual VisitResult visit(VariableDefExpression* expr) override
+	{
+		TokenStreamContext context{expr};
+		auto varExpr = expr->getExpression()->accept(this);
+		return defaultHandler(context, expr, varExpr);
+	}
+
+	virtual VisitResult visit(WithExpression* expr) override
+	{
+		TokenStreamContext context{expr};
+
+		std::vector<VisitResult> variables;
+		for (auto& var : expr->getVariables())
+		{
+			variables.push_back(var->accept(this));
+		}
+
+		auto body = expr->getBody()->accept(this);
+		return defaultHandler(context, expr, variables, body);
+	}
 	/// @}
 
 	/// @name Default handlers
@@ -804,6 +825,60 @@ class ModifyingVisitor : public Visitor
 
 		return {};
 	}
+
+	VisitResult defaultHandler(const TokenStreamContext& context, VariableDefExpression* expr, const VisitResult& exprRet)
+	{
+		if (auto newExpr = std::get_if<Expression::Ptr>(&exprRet))
+		{
+			if (*newExpr)
+				expr->setExpression(*newExpr);
+		}
+		else
+			expr->setExpression(nullptr);
+
+		if (!expr->getExpression())
+			return VisitAction::Delete;
+
+		return {};
+	}
+
+	VisitResult defaultHandler(const TokenStreamContext& context, WithExpression* expr, const std::vector<VisitResult>& varsRet, const VisitResult& bodyRet)
+	{
+		if (auto body = std::get_if<Expression::Ptr>(&bodyRet))
+		{
+			if (*body)
+				expr->setBody(*body);
+		}
+		else
+			expr->setBody(nullptr);
+
+		if (!expr->getBody())
+			return VisitAction::Delete;
+
+		if (std::all_of(varsRet.begin(), varsRet.end(),
+				[](const auto& var) {
+					auto a = std::get_if<Expression::Ptr>(&var);
+					return a && (*a == nullptr);
+				}))
+		{
+			return {};
+		}
+
+		std::vector<Expression::Ptr> newVariables;
+		for (std::size_t i = 0, end = varsRet.size(); i < end; ++i)
+		{
+			if (auto var = std::get_if<Expression::Ptr>(&varsRet[i]))
+			{
+				if (*var)
+					newVariables.push_back(*var);
+				else
+					newVariables.push_back(expr->getVariables()[i]);
+			}
+		}
+
+		expr->setVariables(newVariables);
+		return {};
+	}
 	/// @}
 
 	/**

include/yaramod/utils/observing_visitor.h

@@ -380,6 +380,23 @@ class ObservingVisitor : public Visitor
 	}
 
 	virtual VisitResult visit(RegexpExpression*) override { return {}; }
+
+	virtual VisitResult visit(VariableDefExpression* expr) override
+	{
+		expr->getExpression()->accept(this);
+		return {};
+	}
+
+	virtual VisitResult visit(WithExpression* expr) override
+	{
+		for (const auto& var : expr->getVariables())
+		{
+			var->accept(this);
+		}
+
+		expr->getBody()->accept(this);
+		return {};
+	}
 	/// @}
 
 protected:

include/yaramod/utils/visitor.h

@@ -75,6 +75,8 @@ class ThemExpression;
 class ParenthesesExpression;
 class IntFunctionExpression;
 class RegexpExpression;
+class VariableDefExpression;
+class WithExpression;
 
 /**
  * Abstract class representing visitor design pattern for visiting condition expressions
@@ -149,6 +151,8 @@ class Visitor
 	virtual VisitResult visit(ParenthesesExpression* expr) = 0;
 	virtual VisitResult visit(IntFunctionExpression* expr) = 0;
 	virtual VisitResult visit(RegexpExpression* expr) = 0;
+	virtual VisitResult visit(VariableDefExpression* expr) = 0;
+	virtual VisitResult visit(WithExpression* expr) = 0;
 	/// @}
 	bool resultIsDelete(const VisitResult& result) const
 	{

src/builder/yara_expression_builder.cpp

@@ -1614,4 +1614,55 @@ YaraExpressionBuilder regexp(const std::string& text, const std::string& suffixM
 	return YaraExpressionBuilder(std::move(ts), std::make_shared<RegexpExpression>(std::move(regexp)), Expression::Type::Regexp);
 }
 
+/**
+ * Creates variable definition for "with" expression.
+ *
+ * @param name Name of the token.
+ * @param expr Expression.
+ *
+ * @return Builder.
+ */
+YaraExpressionBuilder var_def(const std::string& name, const YaraExpressionBuilder& expr)
+{
+	std::shared_ptr<TokenStream> ts = std::make_shared<TokenStream>();
+	auto name_token = ts->emplace_back(TokenType::ID, name);
+	ts->emplace_back(TokenType::ASSIGN, "=");
+	ts->moveAppend(expr.getTokenStream());
+	return YaraExpressionBuilder(std::move(ts), std::make_shared<VariableDefExpression>(name_token, expr.get()));
+}
+
+/**
+ * Creates the expression with "with" expression defining variables and body.
+ *
+ * @param vars Variable builders.
+ * @param body Body expression builder.
+ *
+ * @return Builder.
+ */
+YaraExpressionBuilder with(const std::vector<YaraExpressionBuilder>& vars, const YaraExpressionBuilder& body)
+{
+	std::shared_ptr<TokenStream> ts = std::make_shared<TokenStream>();
+	TokenIt with = ts->emplace_back(TokenType::WITH, "with");
+
+	std::vector<Expression::Ptr> varExprs;
+	varExprs.reserve(vars.size());
+
+	for (size_t i = 0; i < vars.size(); ++i)
+	{
+		ts->moveAppend(vars[i].getTokenStream());
+
+		if (i + 1 < vars.size())
+			ts->emplace_back(TokenType::COMMA, ",");
+		else
+			ts->emplace_back(TokenType::COLON, ":");
+
+		varExprs.push_back(vars[i].get());
+	}
+
+	ts->emplace_back(TokenType::LP_WITH_SPACE_AFTER, "(");
+	ts->moveAppend(body.getTokenStream());
+	auto right_bracket = ts->emplace_back(TokenType::RP_WITH_SPACE_BEFORE, ")");
+	return YaraExpressionBuilder(std::move(ts), std::make_shared<WithExpression>(with, std::move(varExprs), body.get(), right_bracket));
+}
+
 }

src/examples/cpp/dump_rules_ast/dumper.h

@@ -587,6 +587,26 @@ class Dumper : public yaramod::ObservingVisitor, public yaramod::ObservingRegexp
 		return {};
 	}
 
+	virtual yaramod::VisitResult visit(yaramod::VariableDefExpression* expr) override
+	{
+		dump("VariableDef", expr, " name=", expr->getName());
+		indentUp();
+		expr->getExpression()->accept(this);
+		indentDown();
+		return {};
+	}
+
+	virtual yaramod::VisitResult visit(yaramod::WithExpression* expr) override
+	{
+		dump("With", expr);
+		indentUp();
+		for (const auto& var : expr->getVariables())
+			var->accept(this);
+		expr->getBody()->accept(this);
+		indentDown();
+		return {};
+	}
+
 	// ==================== ObservingRegexVisitor ====================
 	yaramod::RegexpVisitResult observe(const std::shared_ptr<yaramod::RegexpUnit>& unit)
 	{