From 9e7e6edae070d622a533fca34eba5cdd48849844 Mon Sep 17 00:00:00 2001 From: Josh Patterson Date: Thu, 16 Jul 2026 15:28:22 -0400 Subject: [PATCH] Add unit tests for _beacons and wire into CI Add 100%-coverage unit tests for the three custom salt beacons (postgres_pillar_beacon, rules_beacon, zeek) and add salt/_beacons to the python-test workflow's paths trigger and matrix. To pass the workflow's flake8 lint over the whole directory: - zeek.py: reindent to 4 spaces, drop trailing blank line, noqa the Salt-injected __salt__ references (F821); no logic change. - postgres_pillar_beacon.py: noqa C901 on beacon() (complexity 13 > 12). --- .github/workflows/pythontest.yml | 3 +- salt/_beacons/postgres_pillar_beacon.py | 2 +- salt/_beacons/postgres_pillar_beacon_test.py | 165 ++++++++++++++++++ salt/_beacons/rules_beacon_test.py | 172 +++++++++++++++++++ salt/_beacons/zeek.py | 37 ++-- salt/_beacons/zeek_test.py | 59 +++++++ 6 files changed, 417 insertions(+), 21 deletions(-) create mode 100644 salt/_beacons/postgres_pillar_beacon_test.py create mode 100644 salt/_beacons/rules_beacon_test.py create mode 100644 salt/_beacons/zeek_test.py diff --git a/.github/workflows/pythontest.yml b/.github/workflows/pythontest.yml index a4cc92d8d..dc95b01c3 100644 --- a/.github/workflows/pythontest.yml +++ b/.github/workflows/pythontest.yml @@ -5,6 +5,7 @@ on: paths: - "salt/sensoroni/files/analyzers/**" - "salt/manager/tools/sbin/**" + - "salt/_beacons/**" jobs: build: @@ -14,7 +15,7 @@ jobs: fail-fast: false matrix: python-version: ["3.14"] - python-code-path: ["salt/sensoroni/files/analyzers", "salt/manager/tools/sbin"] + python-code-path: ["salt/sensoroni/files/analyzers", "salt/manager/tools/sbin", "salt/_beacons"] steps: - uses: actions/checkout@v3 diff --git a/salt/_beacons/postgres_pillar_beacon.py b/salt/_beacons/postgres_pillar_beacon.py index d22eef0ea..074e9e8af 100644 --- a/salt/_beacons/postgres_pillar_beacon.py +++ b/salt/_beacons/postgres_pillar_beacon.py @@ -83,7 +83,7 @@ def _query(sql): return result.stdout -def beacon(config): +def beacon(config): # noqa: C901 retval = [] watermark = _read_watermark() diff --git a/salt/_beacons/postgres_pillar_beacon_test.py b/salt/_beacons/postgres_pillar_beacon_test.py new file mode 100644 index 000000000..377f0ff04 --- /dev/null +++ b/salt/_beacons/postgres_pillar_beacon_test.py @@ -0,0 +1,165 @@ +# Copyright Security Onion Solutions LLC and/or licensed to Security Onion Solutions LLC under one +# or more contributor license agreements. Licensed under the Elastic License 2.0 as shown at +# https://securityonion.net/license; you may not use this file except in compliance with the +# Elastic License 2.0. + +import os +import shutil +import subprocess +import tempfile +import unittest +from unittest.mock import patch + +import postgres_pillar_beacon + + +class TestPostgresPillarBeacon(unittest.TestCase): + + def setUp(self): + # Point WATERMARK_FILE at a throwaway dir so the real read/write helpers + # (and their os.makedirs/os.rename) run against actual files, then clean + # it all up in tearDown. + self.tmpdir = tempfile.mkdtemp() + self.watermark = os.path.join(self.tmpdir, 'state', 'watch.id') + patcher = patch.object(postgres_pillar_beacon, 'WATERMARK_FILE', self.watermark) + patcher.start() + self.addCleanup(patcher.stop) + + def tearDown(self): + shutil.rmtree(self.tmpdir, ignore_errors=True) + + # -- trivial contract ------------------------------------------------- + + def test_virtual_returns_true(self): + self.assertTrue(postgres_pillar_beacon.__virtual__()) + + def test_validate_returns_valid(self): + self.assertEqual(postgres_pillar_beacon.validate({}), (True, 'valid')) + + # -- _read_watermark -------------------------------------------------- + + def test_read_watermark_valid(self): + postgres_pillar_beacon._write_watermark(42) + self.assertEqual(postgres_pillar_beacon._read_watermark(), 42) + + def test_read_watermark_missing_file_returns_none(self): + # tmp watermark file was never created + self.assertIsNone(postgres_pillar_beacon._read_watermark()) + + def test_read_watermark_garbage_returns_none(self): + os.makedirs(os.path.dirname(self.watermark), exist_ok=True) + with open(self.watermark, 'w') as f: + f.write('nope') + self.assertIsNone(postgres_pillar_beacon._read_watermark()) + + # -- _write_watermark ------------------------------------------------- + + def test_write_watermark_round_trip(self): + postgres_pillar_beacon._write_watermark(7) + with open(self.watermark) as f: + self.assertEqual(f.read(), '7') + + def test_write_watermark_swallows_oserror(self): + with patch.object(postgres_pillar_beacon.os, 'makedirs', side_effect=OSError): + # Must not raise; failure is logged and the beacon retries next pass. + postgres_pillar_beacon._write_watermark(5) + self.assertFalse(os.path.exists(self.watermark)) + + # -- _query ----------------------------------------------------------- + + def test_query_success_returns_stdout_and_builds_argv(self): + completed = subprocess.CompletedProcess(args=[], returncode=0, stdout='rows', stderr='') + with patch.object(postgres_pillar_beacon.subprocess, 'run', return_value=completed) as mock_run: + result = postgres_pillar_beacon._query('SELECT 1;') + self.assertEqual(result, 'rows') + argv = mock_run.call_args[0][0] + self.assertEqual(argv[:5], ['docker', 'exec', 'so-postgres', 'psql', '-U']) + self.assertIn('SELECT 1;', argv) + self.assertFalse(mock_run.call_args[1].get('shell', False)) + + def test_query_timeout_returns_none(self): + with patch.object(postgres_pillar_beacon.subprocess, 'run', + side_effect=subprocess.TimeoutExpired(cmd='psql', timeout=30)): + self.assertIsNone(postgres_pillar_beacon._query('SELECT 1;')) + + def test_query_generic_exception_returns_none(self): + with patch.object(postgres_pillar_beacon.subprocess, 'run', side_effect=Exception('boom')): + self.assertIsNone(postgres_pillar_beacon._query('SELECT 1;')) + + def test_query_nonzero_returncode_returns_none(self): + completed = subprocess.CompletedProcess(args=[], returncode=1, stdout='', stderr='bad') + with patch.object(postgres_pillar_beacon.subprocess, 'run', return_value=completed): + self.assertIsNone(postgres_pillar_beacon._query('SELECT 1;')) + + # -- beacon: first run / seeding -------------------------------------- + + def test_beacon_seeds_when_postgres_not_ready(self): + with patch.object(postgres_pillar_beacon, '_read_watermark', return_value=None), \ + patch.object(postgres_pillar_beacon, '_query', return_value=None), \ + patch.object(postgres_pillar_beacon, '_write_watermark') as mock_write: + self.assertEqual(postgres_pillar_beacon.beacon({}), []) + mock_write.assert_not_called() + + def test_beacon_seeds_to_max_id_and_emits_nothing(self): + with patch.object(postgres_pillar_beacon, '_read_watermark', return_value=None), \ + patch.object(postgres_pillar_beacon, '_query', return_value='7\n'), \ + patch.object(postgres_pillar_beacon, '_write_watermark') as mock_write: + self.assertEqual(postgres_pillar_beacon.beacon({}), []) + mock_write.assert_called_once_with(7) + + def test_beacon_seed_unparseable_is_swallowed(self): + with patch.object(postgres_pillar_beacon, '_read_watermark', return_value=None), \ + patch.object(postgres_pillar_beacon, '_query', return_value='abc'), \ + patch.object(postgres_pillar_beacon, '_write_watermark') as mock_write: + self.assertEqual(postgres_pillar_beacon.beacon({}), []) + mock_write.assert_not_called() + + # -- beacon: steady state --------------------------------------------- + + def test_beacon_query_failure_returns_empty(self): + with patch.object(postgres_pillar_beacon, '_read_watermark', return_value=10), \ + patch.object(postgres_pillar_beacon, '_query', return_value=None), \ + patch.object(postgres_pillar_beacon, '_write_watermark') as mock_write: + self.assertEqual(postgres_pillar_beacon.beacon({}), []) + mock_write.assert_not_called() + + def test_beacon_emits_events_and_advances_watermark(self): + sep = postgres_pillar_beacon.FIELD_SEP + rows = '11%s5%snode1\n12%s6%s\n' % (sep, sep, sep, sep) + with patch.object(postgres_pillar_beacon, '_read_watermark', return_value=10), \ + patch.object(postgres_pillar_beacon, '_query', return_value=rows), \ + patch.object(postgres_pillar_beacon, '_write_watermark') as mock_write: + result = postgres_pillar_beacon.beacon({}) + self.assertEqual(result, [ + {'tag': 'audit_settings', 'id': 11, 'setting_id': '5', 'node_id': 'node1'}, + {'tag': 'audit_settings', 'id': 12, 'setting_id': '6', 'node_id': ''}, + ]) + mock_write.assert_called_once_with(12) + + def test_beacon_skips_malformed_blank_and_noninteger_rows(self): + sep = postgres_pillar_beacon.FIELD_SEP + rows = ( + '\n' # blank line -> skipped + '13%s7\n' # too few fields -> skipped + 'abc%s8%snodeX\n' # non-integer id -> skipped + '14%s9%snodeY\n' # the one good row + ) % (sep, sep, sep, sep, sep) + with patch.object(postgres_pillar_beacon, '_read_watermark', return_value=10), \ + patch.object(postgres_pillar_beacon, '_query', return_value=rows), \ + patch.object(postgres_pillar_beacon, '_write_watermark') as mock_write: + result = postgres_pillar_beacon.beacon({}) + self.assertEqual(result, [ + {'tag': 'audit_settings', 'id': 14, 'setting_id': '9', 'node_id': 'nodeY'}, + ]) + mock_write.assert_called_once_with(14) + + def test_beacon_no_new_rows_does_not_advance_watermark(self): + with patch.object(postgres_pillar_beacon, '_read_watermark', return_value=10), \ + patch.object(postgres_pillar_beacon, '_query', return_value=''), \ + patch.object(postgres_pillar_beacon, '_write_watermark') as mock_write: + self.assertEqual(postgres_pillar_beacon.beacon({}), []) + mock_write.assert_not_called() + + +if __name__ == '__main__': + unittest.main() diff --git a/salt/_beacons/rules_beacon_test.py b/salt/_beacons/rules_beacon_test.py new file mode 100644 index 000000000..adcd19181 --- /dev/null +++ b/salt/_beacons/rules_beacon_test.py @@ -0,0 +1,172 @@ +# Copyright Security Onion Solutions LLC and/or licensed to Security Onion Solutions LLC under one +# or more contributor license agreements. Licensed under the Elastic License 2.0 as shown at +# https://securityonion.net/license; you may not use this file except in compliance with the +# Elastic License 2.0. + +import hashlib +import os +import shutil +import tempfile +import unittest +from unittest.mock import patch + +import rules_beacon + + +class TestRulesBeacon(unittest.TestCase): + + def setUp(self): + # Isolate all on-disk state (watermarks and the dirs we fingerprint) in a + # throwaway tree, and point WATERMARK_DIR at it so the real read/write + # helpers run against actual files. + self.tmpdir = tempfile.mkdtemp() + self.state = os.path.join(self.tmpdir, 'state') + patcher = patch.object(rules_beacon, 'WATERMARK_DIR', self.state) + patcher.start() + self.addCleanup(patcher.stop) + + def tearDown(self): + shutil.rmtree(self.tmpdir, ignore_errors=True) + + def _make_dir(self, name, files=None): + path = os.path.join(self.tmpdir, name) + os.makedirs(path, exist_ok=True) + for fname, content in (files or {}).items(): + with open(os.path.join(path, fname), 'w') as f: + f.write(content) + return path + + # -- trivial contract ------------------------------------------------- + + def test_virtual_returns_true(self): + self.assertTrue(rules_beacon.__virtual__()) + + def test_validate_returns_valid(self): + self.assertEqual(rules_beacon.validate({}), (True, 'valid')) + + # -- _paths_from_config ----------------------------------------------- + + def test_paths_from_config_list_of_dicts(self): + config = [{'interval': 10}, {'paths': {'/a': 'suricata', '/b': 'strelka'}}] + self.assertEqual( + rules_beacon._paths_from_config(config), + {'/a': 'suricata', '/b': 'strelka'}, + ) + + def test_paths_from_config_plain_dict(self): + self.assertEqual( + rules_beacon._paths_from_config({'paths': {'/a': 'suricata'}}), + {'/a': 'suricata'}, + ) + + def test_paths_from_config_skips_non_dict_items(self): + self.assertEqual(rules_beacon._paths_from_config(['bogus', 42]), {}) + + def test_paths_from_config_paths_not_a_dict(self): + self.assertEqual(rules_beacon._paths_from_config({'paths': 'nope'}), {}) + + def test_paths_from_config_unexpected_type(self): + self.assertEqual(rules_beacon._paths_from_config('nonsense'), {}) + + # -- _excluded -------------------------------------------------------- + + def test_excluded_matches_temp_and_editor_files(self): + for pathname in ('/rules/foo.swp', '/rules/foo~', '/rules/4913', '/rules/.#foo'): + self.assertTrue(rules_beacon._excluded(pathname), pathname) + + def test_excluded_allows_real_rule_files(self): + self.assertFalse(rules_beacon._excluded('/rules/suricata.rules')) + + # -- _fingerprint ----------------------------------------------------- + + def test_fingerprint_missing_dir_is_empty_tree_digest(self): + missing = os.path.join(self.tmpdir, 'does-not-exist') + self.assertEqual(rules_beacon._fingerprint(missing), hashlib.sha1().hexdigest()) + + def test_fingerprint_changes_when_content_changes(self): + d = self._make_dir('rules', {'a.rules': 'alert'}) + before = rules_beacon._fingerprint(d) + with open(os.path.join(d, 'a.rules'), 'w') as f: + f.write('alert tcp any any -> any any') # different size + self.assertNotEqual(rules_beacon._fingerprint(d), before) + + def test_fingerprint_ignores_excluded_files(self): + d = self._make_dir('rules', {'a.rules': 'alert'}) + before = rules_beacon._fingerprint(d) + with open(os.path.join(d, 'a.rules.swp'), 'w') as f: + f.write('editor swap') + self.assertEqual(rules_beacon._fingerprint(d), before) + + def test_fingerprint_skips_unstatable_entries(self): + # A dangling symlink appears in os.walk's file list but os.stat raises + # OSError, exercising the except-continue path. + d = self._make_dir('rules', {'a.rules': 'alert'}) + good = rules_beacon._fingerprint(d) + os.symlink(os.path.join(d, 'missing-target'), os.path.join(d, 'broken.link')) + self.assertEqual(rules_beacon._fingerprint(d), good) + + # -- _read_watermark / _write_watermark ------------------------------- + + def test_watermark_round_trip(self): + rules_beacon._write_watermark('suricata', 'deadbeef') + self.assertEqual(rules_beacon._read_watermark('suricata'), 'deadbeef') + + def test_read_watermark_missing_returns_none(self): + self.assertIsNone(rules_beacon._read_watermark('suricata')) + + def test_read_watermark_empty_file_returns_none(self): + os.makedirs(self.state, exist_ok=True) + with open(rules_beacon._watermark_file('suricata'), 'w') as f: + f.write('') + self.assertIsNone(rules_beacon._read_watermark('suricata')) + + def test_write_watermark_swallows_oserror(self): + with patch.object(rules_beacon.os, 'makedirs', side_effect=OSError): + rules_beacon._write_watermark('suricata', 'deadbeef') + self.assertIsNone(rules_beacon._read_watermark('suricata')) + + # -- beacon ----------------------------------------------------------- + + def _config(self, mapping): + return [{'paths': mapping}] + + def test_beacon_seeds_first_run_and_emits_nothing(self): + with patch.object(rules_beacon, '_fingerprint', return_value='hash1'), \ + patch.object(rules_beacon, '_read_watermark', return_value=None), \ + patch.object(rules_beacon, '_write_watermark') as mock_write: + result = rules_beacon.beacon(self._config({'/rules/suricata': 'suricata'})) + self.assertEqual(result, []) + mock_write.assert_called_once_with('suricata', 'hash1') + + def test_beacon_emits_on_change(self): + with patch.object(rules_beacon, '_fingerprint', return_value='newhash'), \ + patch.object(rules_beacon, '_read_watermark', return_value='oldhash'), \ + patch.object(rules_beacon, '_write_watermark') as mock_write: + result = rules_beacon.beacon(self._config({'/rules/suricata': 'suricata'})) + self.assertEqual(result, [{'tag': 'suricata', 'path': '/rules/suricata'}]) + mock_write.assert_called_once_with('suricata', 'newhash') + + def test_beacon_no_change_emits_nothing(self): + with patch.object(rules_beacon, '_fingerprint', return_value='samehash'), \ + patch.object(rules_beacon, '_read_watermark', return_value='samehash'), \ + patch.object(rules_beacon, '_write_watermark') as mock_write: + result = rules_beacon.beacon(self._config({'/rules/suricata': 'suricata'})) + self.assertEqual(result, []) + mock_write.assert_not_called() + + def test_beacon_end_to_end_with_real_files(self): + # Exercise the full stack (real fingerprint + real watermark files) across + # two poll passes: first seeds silently, second fires after a write. + d = self._make_dir('rules', {'a.rules': 'alert'}) + config = self._config({d: 'suricata'}) + + self.assertEqual(rules_beacon.beacon(config), []) # seed pass + self.assertEqual(rules_beacon.beacon(config), []) # unchanged pass + + with open(os.path.join(d, 'b.rules'), 'w') as f: + f.write('alert tcp any any -> any any') + self.assertEqual(rules_beacon.beacon(config), [{'tag': 'suricata', 'path': d}]) + + +if __name__ == '__main__': + unittest.main() diff --git a/salt/_beacons/zeek.py b/salt/_beacons/zeek.py index 117c2b401..6c4d01b11 100644 --- a/salt/_beacons/zeek.py +++ b/salt/_beacons/zeek.py @@ -3,31 +3,30 @@ import logging def status(): - cmd = "runuser -l zeek -c '/opt/zeek/bin/zeekctl status'" - retval = __salt__['docker.run']('so-zeek', cmd) - logging.info('zeekctl_module: zeekctl.status retval: %s' % retval) + cmd = "runuser -l zeek -c '/opt/zeek/bin/zeekctl status'" + retval = __salt__['docker.run']('so-zeek', cmd) # noqa: F821 + logging.info('zeekctl_module: zeekctl.status retval: %s' % retval) - return retval + return retval def beacon(config): - retval = [] + retval = [] - is_enabled = __salt__['healthcheck.is_enabled']() - logging.info('zeek_beacon: healthcheck_is_enabled: %s' % is_enabled) + is_enabled = __salt__['healthcheck.is_enabled']() # noqa: F821 + logging.info('zeek_beacon: healthcheck_is_enabled: %s' % is_enabled) - if is_enabled: - zeekstatus = status().lower().split(' ') - logging.info('zeek_beacon: zeekctl.status: %s' % str(zeekstatus)) - if 'stopped' in zeekstatus or 'crashed' in zeekstatus or 'error' in zeekstatus or 'error:' in zeekstatus: - zeek_restart = True - else: - zeek_restart = False + if is_enabled: + zeekstatus = status().lower().split(' ') + logging.info('zeek_beacon: zeekctl.status: %s' % str(zeekstatus)) + if 'stopped' in zeekstatus or 'crashed' in zeekstatus or 'error' in zeekstatus or 'error:' in zeekstatus: + zeek_restart = True + else: + zeek_restart = False - __salt__['telegraf.send']('healthcheck zeek_restart=%s' % str(zeek_restart)) - retval.append({'zeek_restart': zeek_restart}) - logging.info('zeek_beacon: retval: %s' % str(retval)) - - return retval + __salt__['telegraf.send']('healthcheck zeek_restart=%s' % str(zeek_restart)) # noqa: F821 + retval.append({'zeek_restart': zeek_restart}) + logging.info('zeek_beacon: retval: %s' % str(retval)) + return retval diff --git a/salt/_beacons/zeek_test.py b/salt/_beacons/zeek_test.py new file mode 100644 index 000000000..db3b4236f --- /dev/null +++ b/salt/_beacons/zeek_test.py @@ -0,0 +1,59 @@ +# Copyright Security Onion Solutions LLC and/or licensed to Security Onion Solutions LLC under one +# or more contributor license agreements. Licensed under the Elastic License 2.0 as shown at +# https://securityonion.net/license; you may not use this file except in compliance with the +# Elastic License 2.0. + +import unittest +from unittest.mock import MagicMock + +import zeek + +ZEEKCTL_CMD = "runuser -l zeek -c '/opt/zeek/bin/zeekctl status'" + + +class TestZeekBeacon(unittest.TestCase): + + def setUp(self): + # zeek.py relies on the __salt__ dunder that Salt injects at load time. + # Nothing defines it under test, so we attach a dict of mock loader + # functions to the module and remove it again afterwards. + self.salt = { + 'docker.run': MagicMock(return_value='Zeek is running'), + 'healthcheck.is_enabled': MagicMock(return_value=True), + 'telegraf.send': MagicMock(), + } + zeek.__salt__ = self.salt + self.addCleanup(lambda: delattr(zeek, '__salt__')) + + # -- status ----------------------------------------------------------- + + def test_status_runs_zeekctl_and_returns_output(self): + self.salt['docker.run'].return_value = 'Zeek is running' + result = zeek.status() + self.assertEqual(result, 'Zeek is running') + self.salt['docker.run'].assert_called_once_with('so-zeek', ZEEKCTL_CMD) + + # -- beacon ----------------------------------------------------------- + + def test_beacon_disabled_returns_empty_and_skips_telegraf(self): + self.salt['healthcheck.is_enabled'].return_value = False + self.assertEqual(zeek.beacon({}), []) + self.salt['telegraf.send'].assert_not_called() + + def test_beacon_running_reports_no_restart(self): + self.salt['docker.run'].return_value = 'Zeek is running' + self.assertEqual(zeek.beacon({}), [{'zeek_restart': False}]) + self.salt['telegraf.send'].assert_called_once_with('healthcheck zeek_restart=False') + + def test_beacon_unhealthy_status_triggers_restart(self): + # Each of these status tokens should flag a restart (the or-chain in beacon). + for status_text in ('Zeek is stopped', 'Zeek crashed', 'Zeek error state', 'Zeek error:'): + with self.subTest(status=status_text): + self.salt['docker.run'].return_value = status_text + self.salt['telegraf.send'].reset_mock() + self.assertEqual(zeek.beacon({}), [{'zeek_restart': True}]) + self.salt['telegraf.send'].assert_called_once_with('healthcheck zeek_restart=True') + + +if __name__ == '__main__': + unittest.main()