mirror of
https://github.com/Security-Onion-Solutions/securityonion.git
synced 2026-07-18 14:53:40 +02:00
Josh Patterson
9e7e6edae0
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).
166 lines
8.0 KiB
Python
166 lines
8.0 KiB
Python
# 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()
|