File Manager - BorneoXploit

Editing: fuzzy_ping.lua

--[[
Copyright (c) 2023, Vsevolod Stakhov <vsevolod@rspamd.com>

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
]]--

local argparse = require "argparse"
local ansicolors = require "ansicolors"
local rspamd_logger = require "rspamd_logger"
local lua_util = require "lua_util"
local lua_fuzzy = require "lua_fuzzy"

local E = {}

local parser = argparse()
    :name 'rspamadm fuzzy_ping'
    :description 'Pings fuzzy storage'
    :help_description_margin(30)
parser:option "-c --config"
      :description "Path to config file"
      :argname("<cfg>")
      :default(rspamd_paths["CONFDIR"] .. "/" .. "rspamd.conf")
parser:option "-r --rule"
      :description "Storage to ping (must be configured in Rspamd configuration)"
      :argname("<name>")
      :default("rspamd.com")
parser:flag "-f --flood"
      :description "Flood mode (no waiting for replies)"
parser:flag "-S --silent"
      :description "Silent mode (statistics only)"
parser:option "-t --timeout"
      :description "Timeout for requests"
      :argname("<timeout>")
      :convert(tonumber)
      :default(5)
parser:option "-s --server"
      :description "Override server to ping"
      :argname("<name>")
parser:option "-n --number"
      :description "Timeout for requests"
      :argname("<number>")
      :convert(tonumber)
      :default(5)
parser:flag "-l --list"
      :description "List configured storages"

local function load_config(opts)
  local _r, err = rspamd_config:load_ucl(opts['config'])

if not _r then
    rspamd_logger.errx('cannot parse %s: %s', opts['config'], err)
    os.exit(1)
  end

-- Init the real structure excluding logging and workers
  _r, err = rspamd_config:parse_rcl({ 'logging', 'worker' })
  if not _r then
    rspamd_logger.errx('cannot process %s: %s', opts['config'], err)
    os.exit(1)
  end

_r, err = rspamd_config:init_modules()
  if not _r then
    rspamd_logger.errx('cannot init modules from %s: %s', opts['config'], err)
    os.exit(1)
  end
end

local function highlight(fmt, ...)
  return ansicolors.white .. string.format(fmt, ...) .. ansicolors.reset
end

local function highlight_err(fmt, ...)
  return ansicolors.red .. string.format(fmt, ...) .. ansicolors.reset
end

local function print_storages(rules)
  for n, rule in pairs(rules) do
    print(highlight('Rule: %s', n))
    print(string.format("\tRead only: %s", rule.read_only))
    -- Handle both unified servers and separate read/write servers
    if rule.servers then
      print(string.format("\tServers: %s", table.concat(lua_util.values(rule.servers), ',')))
    else
      if rule.read_servers then
        print(string.format("\tRead servers: %s", table.concat(lua_util.values(rule.read_servers), ',')))
      end
      if rule.write_servers then
        print(string.format("\tWrite servers: %s", table.concat(lua_util.values(rule.write_servers), ',')))
      end
    end
    print("\tFlags:")

for fl, id in pairs(rule.flags or E) do
      print(string.format("\t\t%s: %s", fl, id))
    end
  end
end

local function std_mean(tbl)
  local function mean()
    local sum = 0
    local count = 0

for _, v in ipairs(tbl) do
      sum = sum + v
      count = count + 1
    end

return (sum / count)
  end

local m
  local vm
  local sum = 0
  local count = 0
  local result

m = mean(tbl)

for _, v in ipairs(tbl) do
    vm = v - m
    sum = sum + (vm * vm)
    count = count + 1
  end

result = math.sqrt(sum / (count - 1))

return result, m
end

local function maxmin(tbl)
  local max = -math.huge
  local min = math.huge

for _, v in ipairs(tbl) do
    max = math.max(max, v)
    min = math.min(min, v)
  end

return max, min
end

local function print_results(results)
  local servers = {}
  local err_servers = {}
  for _, res in ipairs(results) do
    if res.success then
      if servers[res.server] then
        table.insert(servers[res.server], res.latency)
      else
        servers[res.server] = { res.latency }
      end
    else
      if err_servers[res.server] then
        err_servers[res.server] = err_servers[res.server] + 1
      else
        err_servers[res.server] = 1
      end
      -- For the case if no successful replies are detected
      if not servers[res.server] then
        servers[res.server] = {}
      end
    end
  end

for s, l in pairs(servers) do
    local total = #l + (err_servers[s] or 0)
    print(highlight('Summary for %s: %d packets transmitted, %d packets received, %.1f%% packet loss',
        s, total, #l, (total - #l) * 100.0 / total))
    local std, mean = std_mean(l)
    local max, min = maxmin(l)
    print(string.format('round-trip min/avg/max/std-dev = %.2f/%.2f/%.2f/%.2f ms',
        min, mean,
        max, std))
  end
end

local function handler(args)
  local opts = parser:parse(args)

load_config(opts)

-- Perform ping using a fake task from async stuff provided by rspamadm
  local rspamd_task = require "rspamd_task"

-- TODO: this task is not cleared at the end, do something about it some day
  local task = rspamd_task.create(rspamd_config, rspamadm_ev_base)
  task:set_session(rspamadm_session)
  task:set_resolver(rspamadm_dns_resolver)

if opts.list then
    -- The task is only needed to drive the event loop while the upstreams
    -- that are resolved asynchronously (e.g. SRV based) are filled in
    lua_fuzzy.wait_for_storages(rspamd_config, task, nil, opts.timeout, function()
      print_storages(rspamd_plugins.fuzzy_check.list_storages(rspamd_config))
    end)

return
  end

local replied = 0
  local results = {}
  local ping_fuzzy

local function gen_ping_fuzzy_cb(num)
    return function(success, server, latency_or_err)
      if not success then
        if not opts.silent then
          print(highlight_err('error from %s: %s', server, latency_or_err))
        end
        results[num] = {
          success = false,
          error = latency_or_err,
          server = tostring(server),
        }
      else
        if not opts.silent then
          local adjusted_latency = math.floor(latency_or_err * 1000) * 1.0 / 1000;
          print(highlight('reply from %s: %s ms', server, adjusted_latency))

end
        results[num] = {
          success = true,
          latency = latency_or_err,
          server = tostring(server),
        }
      end

if replied == opts.number - 1 then
        print_results(results)
      else
        replied = replied + 1
        if not opts.flood then
          ping_fuzzy(replied + 1)
        end
      end
    end
  end

ping_fuzzy = function(num)
    local ret, err = rspamd_plugins.fuzzy_check.ping_storage(task, gen_ping_fuzzy_cb(num),
        opts.rule, opts.timeout, opts.server)

if not ret then
      print(highlight_err('error from %s: %s', opts.server or opts.rule, err))
      opts.number = opts.number - 1 -- To avoid issues with waiting for other replies
    end
  end

local function start_pings()
    if opts.flood then
      for i = 1, opts.number do
        ping_fuzzy(i)
      end
    else
      ping_fuzzy(1)
    end
  end

if opts.server then
    -- Explicit override, the configured upstreams are not used at all
    start_pings()
  else
    lua_fuzzy.wait_for_storages(rspamd_config, task, { [opts.rule] = true },
        opts.timeout, start_pings)
  end
end

return {
  name = 'fuzzy_ping',
  aliases = { 'fuzzyping' },
  handler = handler,
  description = parser._description
}

Cancel