[Project] Convert corpus_test to new format, document all options

2 files changed, 61 insertions, 23 deletions

diff --git a/lualib/rspamadm/corpus_test.lua b/lualib/rspamadm/corpus_test.lua
index 71d1eeaf1..59b02f649 100644
--- a/lualib/rspamadm/corpus_test.lua
+++ b/lualib/rspamadm/corpus_test.lua
@@ -1,19 +1,48 @@
 local rspamd_logger = require "rspamd_logger"
 local ucl = require "ucl"
 local lua_util = require "lua_util"
-local getopt = require "getopt"
+local argparse = require "argparse"
+
+local parser = argparse()
+    :name "rspamadm corpus_test"
+    :description "Create logs files from email corpus"
+    :help_description_margin(32)
+
+parser:option "-H --ham"
+      :description("Ham directory")
+      :argname("<dir>")
+parser:option "-S --spam"
+      :description("Spam directory")
+      :argname("<dir>")
+parser:option "-n --conns"
+      :description("Number of parallel connections")
+      :argname("<N>")
+      :default(10)
+parser:option "-o --output"
+      :description("Output file")
+      :argname("<file>")
+      :default('results.log')
+parser:option "-t --timeout"
+      :description("Timeout for client connections")
+      :argname("<sec>")
+      :default(60)
+parser:option "-c --connect"
+      :description("Connect to specific host")
+      :argname("<host>")
+      :default('localhost:11334')
+parser:option "-r --rspamc"
+      :description("Use specific rspamc path")
+      :argname("<path>")
+      :default('rspamc')
 
 local HAM = "HAM"
 local SPAM = "SPAM"
 local opts
-local default_opts = {
-  connect = 'localhost:11334',
-}
 
 local function scan_email(n_parallel, path, timeout)
 
-  local rspamc_command = string.format("rspamc --connect %s -j --compact -n %s -t %.3f %s",
-      opts.connect, n_parallel, timeout, path)
+  local rspamc_command = string.format("%s --connect %s -j --compact -n %s -t %.3f %s",
+      opts.rspamc, opts.connect, n_parallel, timeout, path)
   local result = assert(io.popen(rspamc_command))
   result = result:read("*all")
   return result
@@ -24,7 +53,8 @@ local function write_results(results, file)
   local f = io.open(file, 'w')
 
   for _, result in pairs(results) do
-    local log_line = string.format("%s %.2f %s", result.type, result.score, result.action)
+    local log_line = string.format("%s %.2f %s",
+        result.type, result.score, result.action)
 
     for _, sym in pairs(result.symbols) do
       log_line = log_line .. " " .. sym
@@ -96,13 +126,12 @@ local function scan_results_to_logs(results, actual_email_type)
   return logs
 end
 
-return function(args, res)
-  opts = default_opts
-  opts = lua_util.override_defaults(opts, getopt.getopt(args, ''))
-  local ham_directory = res['ham_directory']
-  local spam_directory = res['spam_directory']
-  local connections = res["connections"]
-  local output = res["output_location"]
+local function handler(args)
+  opts = parser:parse(args)
+  local ham_directory = opts['ham_directory']
+  local spam_directory = opts['spam_directory']
+  local connections = opts["connections"]
+  local output = opts["output"]
 
   local results = {}
 
@@ -112,7 +141,7 @@ return function(args, res)
 
   if ham_directory then
     rspamd_logger.messagex("Scanning ham corpus...")
-    local ham_results = scan_email(connections, ham_directory, res["timeout"])
+    local ham_results = scan_email(connections, ham_directory, opts["timeout"])
     ham_results = scan_results_to_logs(ham_results, HAM)
 
     no_of_ham = #ham_results
@@ -124,7 +153,7 @@ return function(args, res)
 
   if spam_directory then
     rspamd_logger.messagex("Scanning spam corpus...")
-    local spam_results = scan_email(connections, spam_directory, res.timeout)
+    local spam_results = scan_email(connections, spam_directory, opts.timeout)
     spam_results = scan_results_to_logs(spam_results, SPAM)
 
     no_of_spam = #spam_results
@@ -145,3 +174,10 @@ return function(args, res)
   rspamd_logger.messagex("No of spam: %s", no_of_spam)
   rspamd_logger.messagex("Messages/sec: %s", (total_msgs / elapsed_time))
 end
+
+
+return {
+  name = 'corpus_test',
+  handler = handler,
+  description = parser._description
+}
\ No newline at end of file
diff --git a/lualib/rspamadm/grep.lua b/lualib/rspamadm/grep.lua
index f20ac76e5..0af83c1cf 100644
--- a/lualib/rspamadm/grep.lua
+++ b/lualib/rspamadm/grep.lua
@@ -22,15 +22,17 @@ local parser = argparse()
     :name "rspamadm grep"
     :description "Search for patterns in rspamd logs"
     :help_description_margin(30)
-parser:option "-s --string"
-      :description('Plain string to search (case-insensitive)')
-      :argname "<str>"
+parser:mutex(
+    parser:option "-s --string"
+          :description('Plain string to search (case-insensitive)')
+          :argname "<str>",
+    parser:option "-p --pattern"
+          :description('Pattern to search for (regex)')
+          :argname "<re>"
+)
 parser:flag "-l --lua"
       :description('Use Lua patterns in string search')
-parser:option "-p --pattern"
-      :description('Pattern to search for (regex)')
-      :args(1)
-      :argname "<re>"
+
 parser:argument "input":args "*"
       :description('Process specified inputs')
       :default("stdin")