From 404a439fef09567b522a4dc9fa9cb19dbad94bf7 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Mon, 13 Feb 2017 12:52:23 +0000
Subject: [PATCH] [Feature] Implement received headers flags

---
 rules/global_functions.lua | 24 ++++++++++++++----------
 src/libmime/message.c      |  1 +
 src/libmime/message.h      |  5 +++++
 src/libmime/mime_headers.c | 10 ++++++++++
 src/lua/lua_task.c         | 34 +++++++++++++++++++++++++++++++++-
 5 files changed, 63 insertions(+), 11 deletions(-)

diff --git a/rules/global_functions.lua b/rules/global_functions.lua
index f98587b47..439854f28 100644
--- a/rules/global_functions.lua
+++ b/rules/global_functions.lua
@@ -307,17 +307,19 @@ local function meta_received_function(task)
   local invalid_factor = 0
   local rh = task:get_received_headers()
   local time_factor = 0
-  local tls_factor = 0
+  local secure_factor = 0
+  local fun = require "fun"
 
   if rh and #rh > 0 then
 
-    local ntotal = 1.0 * #rh
-    count_factor = 1.0 / ntotal
+    local ntotal = 0.0
     local init_time = 0
 
-    for _,rc in ipairs(rh) do
+    fun.each(function(rc)
+      ntotal = ntotal + 1.0
+
       if not rc.by_hostname then
-        invalid_factor = invalid_factor + 1
+        invalid_factor = invalid_factor + 1.0
       end
       if init_time == 0 and rc.timestamp then
         init_time = rc.timestamp
@@ -325,20 +327,22 @@ local function meta_received_function(task)
         time_factor = time_factor + math.abs(init_time - rc.timestamp)
         init_time = rc.timestamp
       end
-      if rc.proto and (rc.proto == 'esmtps') then
-        tls_factor = tls_factor + 1
+      if rc.flags and (rc.flags['ssl'] or rc.flags['authenticated']) then
+        secure_factor = secure_factor + 1.0
       end
-    end
+    end,
+    fun.filter(function(rc) return not rc.flags or not rc.flags['artificial'] end, rh))
 
     invalid_factor = invalid_factor / ntotal
-    tls_factor = tls_factor / ntotal
+    secure_factor = secure_factor / ntotal
+    count_factor = 1.0 / ntotal
 
     if time_factor ~= 0 then
       time_factor = 1.0 / time_factor
     end
   end
 
-  return {count_factor, invalid_factor, time_factor, tls_factor}
+  return {count_factor, invalid_factor, time_factor, secure_factor}
 end
 
 local function meta_urls_function(task)
diff --git a/src/libmime/message.c b/src/libmime/message.c
index 6692c76e2..99384e8a2 100644
--- a/src/libmime/message.c
+++ b/src/libmime/message.c
@@ -725,6 +725,7 @@ rspamd_message_parse (struct rspamd_task *task)
 
 			trecv = rspamd_mempool_alloc0 (task->task_pool,
 					sizeof (struct received_header));
+			trecv->flags |= RSPAMD_RECEIVED_FLAG_ARTIFICIAL;
 			trecv->real_ip = rspamd_mempool_strdup (task->task_pool,
 					rspamd_inet_address_to_string (task->from_addr));
 			trecv->from_ip = trecv->real_ip;
diff --git a/src/libmime/message.h b/src/libmime/message.h
index da816ad43..8c0f919ea 100644
--- a/src/libmime/message.h
+++ b/src/libmime/message.h
@@ -103,6 +103,10 @@ enum rspamd_received_type {
 	RSPAMD_RECEIVED_UNKNOWN
 };
 
+#define RSPAMD_RECEIVED_FLAG_ARTIFICIAL (1 << 0)
+#define RSPAMD_RECEIVED_FLAG_SSL (1 << 1)
+#define RSPAMD_RECEIVED_FLAG_AUTHENTICATED (1 << 2)
+
 struct received_header {
 	gchar *from_hostname;
 	gchar *from_ip;
@@ -114,6 +118,7 @@ struct received_header {
 	struct rspamd_mime_header *hdr;
 	time_t timestamp;
 	enum rspamd_received_type type;
+	gint flags;
 };
 
 /**
diff --git a/src/libmime/mime_headers.c b/src/libmime/mime_headers.c
index 3cac797d6..75ec1b96d 100644
--- a/src/libmime/mime_headers.c
+++ b/src/libmime/mime_headers.c
@@ -40,6 +40,16 @@ rspamd_mime_header_check_special (struct rspamd_task *task,
 		recv->hdr = rh;
 		rspamd_smtp_recieved_parse (task, rh->decoded,
 				strlen (rh->decoded), recv);
+		/* Set flags */
+		if (recv->type == RSPAMD_RECEIVED_ESMTPA ||
+				recv->type == RSPAMD_RECEIVED_ESMTPSA) {
+			recv->flags |= RSPAMD_RECEIVED_FLAG_AUTHENTICATED;
+		}
+		if (recv->type == RSPAMD_RECEIVED_ESMTPS ||
+				recv->type == RSPAMD_RECEIVED_ESMTPSA) {
+			recv->flags |= RSPAMD_RECEIVED_FLAG_SSL;
+		}
+
 		g_ptr_array_add (task->received, recv);
 		break;
 	case 0x76F31A09F4352521ULL:	/* to */
diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index 7ae0bd99d..9b15eb2c0 100644
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -1619,12 +1619,44 @@ lua_task_get_received_headers (lua_State * L)
 			for (i = 0; i < task->received->len; i ++) {
 				rh = g_ptr_array_index (task->received, i);
 
-				lua_createtable (L, 0, 9);
+				lua_createtable (L, 0, 10);
 
 				if (rh->hdr && rh->hdr->decoded) {
 					rspamd_lua_table_set (L, "raw", rh->hdr->decoded);
 				}
 
+				lua_pushstring (L, "flags");
+				lua_createtable (L, 0, 3);
+
+				lua_pushstring (L, "artificial");
+				if (rh->flags & RSPAMD_RECEIVED_FLAG_ARTIFICIAL) {
+					lua_pushboolean (L, true);
+				}
+				else {
+					lua_pushboolean (L, false);
+				}
+				lua_settable (L, -3);
+
+				lua_pushstring (L, "authenticated");
+				if (rh->flags & RSPAMD_RECEIVED_FLAG_AUTHENTICATED) {
+					lua_pushboolean (L, true);
+				}
+				else {
+					lua_pushboolean (L, false);
+				}
+				lua_settable (L, -3);
+
+				lua_pushstring (L, "ssl");
+				if (rh->flags & RSPAMD_RECEIVED_FLAG_SSL) {
+					lua_pushboolean (L, true);
+				}
+				else {
+					lua_pushboolean (L, false);
+				}
+				lua_settable (L, -3);
+
+				lua_settable (L, -3);
+
 				if (G_UNLIKELY (rh->from_ip == NULL &&
 						rh->real_ip == NULL &&
 						rh->real_hostname == NULL &&
-- 
2.39.5