munin_rspamd/rspamd_

#!/usr/bin/env python3
# %# family=auto
# %# capabilities=autoconf suggest

"""
Rpsamd munin plugin

Can be use as
- rspamd_actions
- rspamd_results
- rspamd_scantime
- rspamd_uptime
- rspamd_memory_usage
- rspamd_memory_stats
- rspamd_connections
- rspamd_ctrl_connections
- rspamd_statfiles
- rspamd_fuzzy_hashes

Config :

[rspamd_*]
env.rspamc = "/usr/bin/rspamc"
env.cache_file = "/tmp/munin_rspamd.cache"
env.cache_time = 240
"""

import json
import os
import re
import shlex
import sys
import time
import io
import unittest
import unittest.mock
import tempfile
import contextlib


def usage():
    """Return simple help"""
    return """You should ln -s /usr/share/munin/plugins/rspamd_ /etc/munin/plugins/rspamd_FUNC
With FUNC one of %s""" % (
        ", ".join(CONFIGS.keys())
    )


def avg_scantime(scantimes):
    """Returns average, min, max scantimes"""
    total = [float(t) for t in scantimes if t]
    return sum(total) / len(total), min(total), max(total)


def storage2name(name):
    """Format storage name"""
    return re.sub(r"[\.\- ]", "_", name)


CONFIGS = {
    "actions": [
        lambda data: """\
graph_title Rspamd actions
graph_vlabel mails
graph_args --base 1000 -l 0
graph_category antispam
graph_scale yes
scanned.label Scanned/total
learned.label Learned
%s"""
        % (
            "\n".join(
                [
                    "%s.label %s" % (action.replace(" ", "_"), action.title())
                    for action in data["actions"]
                ]
            )
        ),
        lambda data: """\
scanned.value %d
learned.value %d
%s"""
        % (
            data["scanned"],
            data["learned"],
            "\n".join(
                [
                    "%s.value %d" % (action.replace(" ", "_"), value)
                    for action, value in data["actions"].items()
                ]
            ),
        ),
    ],
    "results": [
        """\
graph_title Rspamd results
graph_vlabel mails
graph_args --base 1000 -l 0
graph_category antispam
graph_scale yes
spam.label Spams
spam.draw AREA
ham.label Ham
ham.draw STACK""",
        lambda data: """\
spam.value %d
ham.value %d"""
        % (data["spam_count"], data["ham_count"]),
    ],
    "scantime": [
        """\
graph_title Rspamd scantime
graph_args --base 1000 -l 0
graph_category antispam
graph_scale yes
graph_vlabel seconds
scantime.label Average scantime
min.label Minimum scantime
max.label Maximum scantime
""",
        lambda data: """\
scantime.value %f
min.value %f
max.value %f"""
        % avg_scantime(data["scan_times"]),
    ],
    "uptime": [
        """\
graph_title Rspamd uptime
graph_args --base 1000 -l 0
graph_category antispam
graph_scale no
graph_vlabel uptime in days
uptime.label uptime
uptime.draw AREA
""",
        lambda data: "uptime.value %f" % (data["uptime"] / (3600 * 24)),
    ],
    "memory_usage": [
        """\
graph_title Rspamd memory usage
graph_args --base 1024 -l 0
graph_category antispam
graph_scale yes
graph_vlabel bytes
memory.label Memory usage
""",
        lambda data: "memory.value %d" % data["bytes_allocated"],
    ],
    "memory_stats": [
        """\
graph_title Rspamd memory stats
graph_args --base 1000 -l 0
graph_category antispam
graph_scale yes
pools.label Pools allocated - freed
chunks.label Chunks allocated - freed
shared_chunks_allocated.label Shared chunks allocated
""",
        lambda data: """\
pools.value %d
chunks.value %d
shared_chunks_allocated.value %d"""
        % (
            data["pools_allocated"] - data["pools_freed"],
            data["chunks_allocated"] - data["chunks_freed"],
            data["shared_chunks_allocated"],
        ),
    ],
    "connections": [
        """\
graph_title Rspamd connections
graph_args --base 1000 -l 0
graph_category antispam
graph_scale yes
graph_vlabel connections
connections.label opened connections
""",
        lambda data: "connections.value %d" % data["connections"],
    ],
    "ctrl_connections": [
        """\
graph_title Rspamd control connections count
graph_args --base 1000 -l 0
graph_category antispam
graph_scale yes
graph_vlabel connections
connections.label control connections count
""",
        lambda data: "connections.value %d" % data["control_connections"],
    ],
    "statfiles": [
        lambda data: """\
graph_title rspamd bayes statfiles
graph_args --base 1000 -l 0
graph_scale yes
graph_category antispam
graph_vlabel learned mails
%s"""
        % (
            "\n".join(
                [
                    """learned_{sym}.label {sym} learned
users_{sym}.label {sym} users""".format(
                        sym=stat["symbol"]
                    )
                    for stat in data["statfiles"]
                ]
            )
        ),
        lambda data: "\n".join(
            """\
learned_{sym}.value {lv}
users_{sym}.value {uv}""".format(
                sym=stat["symbol"], lv=stat["revision"], uv=stat["users"]
            )
            for stat in data["statfiles"]
        ),
    ],
    "fuzzy_hashes": [
        lambda data: """\
graph_title rspamd fuzzy hases storage
graph_args --base 1000 -l 0
graph_scale yes
graph_category antispam
graph_vlabel hashes
%s"""
        % (
            "\n".join(
                [
                    "%s.label %s" % (storage2name(storage), storage)
                    for storage in data["fuzzy_hashes"]
                ]
            )
        ),
        lambda data: "\n".join(
            [
                "%s.value %d" % (storage2name(storage), count)
                for storage, count in data["fuzzy_hashes"].items()
            ]
        ),
    ],
}


def main(argv, environ):
    """main function"""
    rspamc = environ.get("rspamc", "/usr/bin/rspamc")
    cache_file = environ.get("cache_file", "/tmp/munin_rspamd.cache")
    cache_time = int(environ.get("cache_time", 60 * 4))

    rspamc_cmd = shlex.join([rspamc, "--json", "stat"])

    if len(argv) > 1:
        if argv[1] == "autoconf":
            if os.path.exists(rspamc):
                return "yes"
            return "no (%s not found)" % rspamc
        if argv[1] == "suggest":
            return "#./rspamd_ suggest\n" + ("\n".join(CONFIGS.keys()))

    data = None
    if os.path.exists(cache_file):
        cache_stat = os.stat(cache_file)
        if cache_stat.st_mtime > time.time() - cache_time:
            with open(cache_file, "r", encoding="utf-8") as cache:
                try:
                    data = json.loads(cache.read())
                except json.JSONDecodeError:
                    pass

    if data is None:
        with open(cache_file, "w", encoding="utf-8") as cache:
            with os.popen(rspamc_cmd, mode="r") as pipe:
                data = pipe.read()
                cache.write(data)
                data = json.loads(data)

    spl = argv[0].split("_")[1:]
    if len(spl) == 0:
        raise NameError(usage())
    cmd = "_".join(spl)
    if cmd not in CONFIGS:
        raise NameError("Unknown command %r. %s" % (cmd, usage()))

    config = CONFIGS[cmd][0]
    act = CONFIGS[cmd][1]

    if len(argv) > 1 and argv[1] == "config":
        if not isinstance(config, str):
            config = config(data)
        return config

    return act(data)


if __name__ == "__main__":
    try:
        print(main(sys.argv, os.environ))
    except NameError as main_expt:
        print(main_expt)

#
# Tests
#

JSON = r"""{\
"version":"3.4","config_id":"1234","uptime":16991,"read_only":false,\
"scanned":43,"learned":43,"actions":\
{"reject":0,"soft reject":0,"rewrite subject":0,"add header":5,\
"greylist":4,"no action":34},"scan_times":\
[0.910916,1.291352,1.241323,null,null,null,null,null,null,null,\
null,null,null,null,null,null,null,null,null,null,null,null,null\
,null,null,null,null,null,null,null,null],\
"spam_count":5,"ham_count":38,"connections":2,\
"control_connections":265,"pools_allocated":538,"pools_freed":512,\
"bytes_allocated":27817136,"chunks_allocated":110,\
"shared_chunks_allocated":3,"chunks_freed":0,"chunks_oversized":1,\
"fragmented":0,"total_learns":40,"statfiles":\
[{"revision":33,"used":0,"total":0,"size":0,"symbol":"BAYES_SPAM",\
"type":"redis","languages":0,"users":1},\
{"revision":7,"used":0,"total":0,"size":0,"symbol":"BAYES_HAM",\
"type":"redis","languages":0,"users":1}],\
"fuzzy_hashes":{"local":0,"rspamd.com":1872513971},"scan_time":0.02}"""

EXPTS = {
    "actions": [
        "reject.value 0",
        "soft_reject.value 0",
        "rewrite_subject.value 0",
        "add_header.value 5",
        "greylist.value 4",
        "no_action.value 34",
        "learned.value 43",
        "scanned.value 43",
    ],
    "results": ["spam.value 5", "ham.value 38"],
    "scantime": ["scantime.value 1.147864", "min.value 0.910916", "max.value 1.291352"],
    "uptime": [
        "uptime.value 0.196655",
    ],
    "memory_usage": [
        "memory.value 27817136",
    ],
    "memory_stats": [
        "chunks.value 110",
        "shared_chunks_allocated.value 3",
        "pools.value 26",
    ],
    "connections": ["connections.value 2"],
    "ctrl_connections": ["connections.value 265"],
    "statfiles": [
        "users_BAYES_SPAM.value 1",
        "learned_BAYES_HAM.value 7",
        "learned_BAYES_SPAM.value 33",
        "users_BAYES_HAM.value 1",
    ],
    "fuzzy_hashes": ["local.value 0", "rspamd_com.value 1872513971"],
}


@contextlib.contextmanager
def json_sample():
    """Mock rspamc --json stat result"""
    return io.StringIO(JSON)


class MuninRspamdTests(unittest.TestCase):
    """Plugin selftest"""

    def setUp(self):
        """Prepare cache file & environment"""
        self.json = lambda: io.StringIO(JSON)
        tmpfd, self.tmpcache = tempfile.mkstemp()
        os.close(tmpfd)
        os.unlink(self.tmpcache)
        self.environ = {"cache_file": self.tmpcache}

    def tearDown(self):
        """Remove cache file"""
        if os.path.isfile(self.tmpcache):
            os.unlink(self.tmpcache)

    def test_configs(self):
        """check plugin config runs without error and returns
        expected fields
        """
        for cmd in CONFIGS:
            with self.subTest(cmd=cmd):
                with unittest.mock.patch("os.popen", return_value=self.json()):
                    ret = main(["rspamd_%s" % cmd, "config"], self.environ)
                    self.assertIn("graph_title", ret)
                    self.assertIn("graph_args", ret)
                    self.assertIn("graph_scale", ret)
                    self.assertIn("graph_category antispam", ret)
                self.tearDown()

    def test_values(self):
        """check plugin runs without error"""
        for cmd in CONFIGS:
            with self.subTest(cmd=cmd):
                with unittest.mock.patch("os.popen", return_value=self.json()):
                    ret = main(["rspamd_%s" % cmd], self.environ)
                    self.assertNotEqual(len(ret), 0)

    def test_consistency(self):
        """check config & values matches"""
        for cmd in CONFIGS:
            with self.subTest(cmd=cmd):
                with unittest.mock.patch("os.popen", return_value=self.json()):
                    conf = main([f"rspamd_{cmd}", "config"], self.environ)
                    values = main([f"rspamd_{cmd}"], self.environ)

                    keys = []
                    for line in conf.split("\n"):
                        prefix = line.split(" ")[0]
                        if "." not in prefix:
                            continue
                        spl = prefix.split(".")
                        if len(spl) < 2 or spl[1] != "label":
                            continue
                        keys.append(spl[0])

                    for name in keys:
                        self.assertIn(f"{name}.value ", values)

    def test_results(self):
        """munin values"""
        for cmd, expt in EXPTS.items():
            with self.subTest(cmd=cmd):
                with unittest.mock.patch("os.popen", return_value=self.json()):
                    ret = main([f"rspamd_{cmd}"], self.environ)
                    self.assertEqual(set(ret.split("\n")), set(expt))

    def test_suggest(self):
        """munin suggest autoconf test"""
        ret = main(["rspamd_", "suggest"], self.environ)
        self.assertEqual(
            set(ret.split("\n")), set(CONFIGS.keys()) | set(["#./rspamd_ suggest"])
        )