From e29277ca2048c114459ca7c6c6b25b6e1e0f2ee1 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Tue, 12 Jun 2018 15:42:58 +0100
Subject: [PATCH] [Feature] Add maps based on Top Level Domains

---
 src/plugins/lua/multimap.lua             | 6 ++++++
 test/functional/cases/102_multimap.robot | 8 ++++++++
 test/functional/configs/maps/top.list    | 2 ++
 test/functional/configs/multimap.conf    | 5 +++++
 4 files changed, 21 insertions(+)
 create mode 100644 test/functional/configs/maps/top.list

diff --git a/src/plugins/lua/multimap.lua b/src/plugins/lua/multimap.lua
index 7e83988c6..8a2b21ddf 100644
--- a/src/plugins/lua/multimap.lua
+++ b/src/plugins/lua/multimap.lua
@@ -93,6 +93,9 @@ local function apply_hostname_filter(task, filter, hostname, r)
   if filter == 'tld' then
     local tld = util.get_tld(hostname)
     return tld
+  elseif filter == 'top' then
+    local tld = util.get_tld(hostname)
+    return tld:match('[^.]*$') or tld
   else
     if not r['re_filter'] then
       local pat = string.match(filter, 'tld:regexp:(.+)')
@@ -123,6 +126,9 @@ local function apply_url_filter(task, filter, url, r)
 
   if filter == 'tld' then
     return url:get_tld()
+  elseif filter == 'top' then
+    local tld = url:get_tld()
+    return tld:match('[^.]*$') or tld
   elseif filter == 'full' then
     return url:get_text()
   elseif filter == 'is_phished' then
diff --git a/test/functional/cases/102_multimap.robot b/test/functional/cases/102_multimap.robot
index c46fff39b..7aac43e4e 100644
--- a/test/functional/cases/102_multimap.robot
+++ b/test/functional/cases/102_multimap.robot
@@ -133,6 +133,14 @@ MAP - HOSTNAME MISS
   ${result} =  Scan Message With Rspamc  ${MESSAGE}  --ip  127.0.0.1  --hostname  rspamd.com
   Check Rspamc  ${result}  HOSTNAME_MAP  inverse=1
 
+MAP - TOP
+  ${result} =  Scan Message With Rspamc  ${MESSAGE}  --ip  127.0.0.1  --hostname  example.com.au
+  Check Rspamc  ${result}  HOSTNAME_TOP_MAP
+
+MAP - TOP MISS
+  ${result} =  Scan Message With Rspamc  ${MESSAGE}  --ip  127.0.0.1  --hostname  example.com.bg
+  Check Rspamc  ${result}  HOSTNAME_TOP_MAP  inverse=1
+
 MAP - CDB - HOSTNAME
   ${result} =  Scan Message With Rspamc  ${MESSAGE}  --ip  127.0.0.1  --hostname  example.com
   Check Rspamc  ${result}  CDB_HOSTNAME
diff --git a/test/functional/configs/maps/top.list b/test/functional/configs/maps/top.list
new file mode 100644
index 000000000..d8a152ed8
--- /dev/null
+++ b/test/functional/configs/maps/top.list
@@ -0,0 +1,2 @@
+au
+#bg
\ No newline at end of file
diff --git a/test/functional/configs/multimap.conf b/test/functional/configs/multimap.conf
index 98701794f..61bd310a8 100644
--- a/test/functional/configs/multimap.conf
+++ b/test/functional/configs/multimap.conf
@@ -58,6 +58,11 @@ multimap {
     type = "hostname";
     map = "${TESTDIR}/configs/maps/domains.list";
   }
+  HOSTNAME_TOP_MAP {
+    type = "hostname";
+    filter = "top";
+    map = "${TESTDIR}/configs/maps/top.list";
+  }
   CDB_HOSTNAME {
     type = "hostname";
     map = "cdb://${TESTDIR}/configs/maps/domains.cdb";
-- 
2.39.5