def test_add_logger(self, socket_log_mock, mkdirs_mock):
mngr_conf = {'root_log_directory': '/totally/a/../fake/dir/../name'}
# We'll use this to make sure directory paths are cleaned
cleaned_dir_path = os.path.normpath(mngr_conf['root_log_directory'])
lm = bsc.StreamCaptureManager(mngr_conf, [])
lm.add_logger('foo', ['', 9000], 'udp', '/tmp')
assert len(lm._stream_capturers.keys()) == 1
assert "['', 9000]" in lm._stream_capturers
# Default root_log_directory usage and normalization check
lm.add_logger('baz', ['', 8500], 'udp')
socket_log_mock.assert_called_with(
{
'log_dir': cleaned_dir_path,
'name': 'baz',
'rotate_log': True,
'pre_write_transforms': [],
}, ['', 8500], 'udp')
assert lm._pool.free_count() == 48
# Check to make sure that home directory expansion is being done
socket_log_mock.reset_mock()
lm.add_logger('testlog', ['', 1234], 'udp', '~/logger_dir')
expanded_user_path = os.path.expanduser('~/logger_dir')
socket_log_mock.assert_called_with(
{
'log_dir': expanded_user_path,
'name': 'testlog',
'rotate_log': True,
'pre_write_transforms': [],
}, ['', 1234], 'udp')
def test_log_manager_init(self, socket_log_mock):
loggers = [
('foo', ['', 9000], 'udp', '/tmp', {'rotate_log': True}),
('bar', ['', 8125], 'udp', '/tmp', {}),
]
fake_mngr_conf = 'mngr_conf'
lm = bsc.StreamCaptureManager(fake_mngr_conf, loggers)
assert lm._mngr_conf == fake_mngr_conf
assert len(lm._stream_capturers.keys()) == 2
assert "['', 9000]" in lm._stream_capturers.keys()
assert "['', 8125]" in lm._stream_capturers.keys()
def main():
ap = argparse.ArgumentParser(
description=__doc__,
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
args = ap.parse_args() # noqa
if not os.path.isfile(config_file):
print(
"Unable to locate config. Starting up handlers with default values ..."
)
host = "localhost"
port = "8080"
handler_configs = []
root_log_dir = "/tmp"
mngr_conf = {"root_log_directory": root_log_dir}
else:
with open(config_file) as log_conf:
conf = yaml.load(log_conf, Loader=yaml.Loader)
mngr_conf = conf["capture_manager"]
host = mngr_conf["manager_server"]["host"]
port = mngr_conf["manager_server"]["port"]
handler_configs = []
for handler_conf in conf["handlers"]:
if "path" in handler_conf:
handler_path = handler_conf.pop("path")
if not os.path.isabs(handler_path):
handler_path = os.path.join(
mngr_conf["root_log_directory"], handler_path)
else:
handler_path = mngr_conf["root_log_directory"]
handler_configs.append((
handler_conf.pop("name"),
handler_conf.pop("address"),
handler_conf.pop("conn_type"),
handler_path,
handler_conf,
))
lgr_mngr = bsc.StreamCaptureManager(mngr_conf, handler_configs)
manager_server = bsc.StreamCaptureManagerServer(logger_manager=lgr_mngr,
host=host,
port=port)
t = threading.Thread(target=manager_server.start)
t.setDaemon(True)
t.start()
lgr_mngr.run_socket_event_loop()
def test_forced_log_rotation(self, socket_mock, pcap_open_mock):
''''''
lm = bsc.StreamCaptureManager(None, [])
with mock.patch('os.mkdir') as mkdir_mock:
lm.add_logger('foo', ['', 9000], 'udp', '/tmp')
lm.add_logger('bar', ['', 8500], 'udp', '/tmp')
bar = lm._stream_capturers["['', 8500]"][0]
pre_rot_count = pcap_open_mock.call_count
lm.rotate_capture_handler_log('bar')
post_rot_count = pcap_open_mock.call_count
assert post_rot_count - pre_rot_count == 1
def test_forced_log_rotation(self, socket_mock, pcap_open_mock):
""""""
lm = bsc.StreamCaptureManager(None, [])
with mock.patch("os.mkdir") as mkdir_mock:
lm.add_logger("foo", ["", 9000], "udp", "/tmp")
lm.add_logger("bar", ["", 8500], "udp", "/tmp")
bar = lm._stream_capturers["['', 8500]"][0]
pre_rot_count = pcap_open_mock.call_count
lm.rotate_capture_handler_log("bar")
post_rot_count = pcap_open_mock.call_count
assert post_rot_count - pre_rot_count == 1
def test_get_logger_stats(self, socket_log_mock):
lm = bsc.StreamCaptureManager(None, [])
with mock.patch('os.mkdir') as mkdir_mock:
lm.add_logger('foo', ['', 9000], 'udp', '/tmp')
lm.add_logger('bar', ['', 8500], 'udp', '/tmp')
logger_data = lm.get_handler_stats()
# Note we're not going to test content of the returned data because
# that is handled by SocketStreamCapturer. There is an appropriate test
# for that in the SocketStreamCapturer section.
assert len(logger_data.keys()) == 2
assert "['', 8500]" in logger_data.keys()
assert "['', 9000]" in logger_data.keys()
def test_get_logger_data(self, socket_log_mock):
lm = bsc.StreamCaptureManager(None, [])
with mock.patch("os.mkdir") as mkdir_mock:
lm.add_logger("foo", ["", 9000], "udp", "/tmp")
lm.add_logger("bar", ["", 8500], "udp", "/tmp")
logger_data = lm.get_logger_data()
# Note we're not going to test content of the returned data because
# that is handled by SocketStreamCapturer. There is an appropriate test
# for that in the SocketStreamCapturer section.
assert len(logger_data.keys()) == 2
assert "['', 8500]" in logger_data.keys()
assert "['', 9000]" in logger_data.keys()
def main():
ap = argparse.ArgumentParser(
description = __doc__,
formatter_class = argparse.ArgumentDefaultsHelpFormatter
)
args = ap.parse_args()
if not os.path.isfile(config_file):
print('Unable to locate config. Starting up handlers with default values ...')
host = 'localhost'
port = '8080'
handler_configs = []
root_log_dir = '/tmp'
mngr_conf = {
'root_log_directory': root_log_dir
}
else:
with open(config_file) as log_conf:
conf = yaml.load(log_conf, Loader=yaml.Loader)
mngr_conf = conf['capture_manager']
host = mngr_conf['manager_server']['host']
port = mngr_conf['manager_server']['port']
handler_configs = []
for handler_conf in conf['handlers']:
if 'path' in handler_conf:
handler_path = handler_conf.pop('path')
if not os.path.isabs(handler_path):
handler_path = os.path.join(mngr_conf['root_log_directory'], handler_path)
else:
handler_path = mngr_conf['root_log_directory']
handler_configs.append((
handler_conf.pop('name'),
handler_conf.pop('address'),
handler_conf.pop('conn_type'),
handler_path,
handler_conf
))
lgr_mngr = bsc.StreamCaptureManager(mngr_conf, handler_configs)
manager_server = bsc.StreamCaptureManagerServer(logger_manager=lgr_mngr, host=host, port=port)
t = threading.Thread(target=manager_server.start)
t.setDaemon(True)
t.start()
lgr_mngr.run_socket_event_loop()
def test_log_manager_init(self, socket_log_mock):
loggers = [
("foo", ["", 9000], "udp", "/tmp", {
"rotate_log": True
}),
("bar", ["", 8125], "udp", "/tmp", {}),
]
fake_mngr_conf = "mngr_conf"
lm = bsc.StreamCaptureManager(fake_mngr_conf, loggers)
assert lm._mngr_conf == fake_mngr_conf
assert len(lm._stream_capturers.keys()) == 2
assert "['', 9000]" in lm._stream_capturers.keys()
assert "['', 8125]" in lm._stream_capturers.keys()
def test_bad_type_transform_load(self, socket_mock, mkdirs_mock, open_mock,
log_mock):
logging.getLogger('ait').setLevel(logging.INFO)
mngr_conf = {'root_log_directory': '/tmp'}
lm = bsc.StreamCaptureManager(mngr_conf, [])
bad_func_name = ('foobarbaz', )
kwargs = {'pre_write_transforms': [bad_func_name]}
lm.add_logger('testlog', ['', 9876], 'udp', '~/logger_dir', **kwargs)
msg = 'Unable to determine how to load data transform "{}"'.format(
bad_func_name)
log_mock.assert_called_with(msg)
logging.getLogger('ait').setLevel(logging.CRITICAL)
def test_bad_builtin_transform_load(self, socket_mock, mkdirs_mock,
open_mock, log_mock):
logging.getLogger('ait').setLevel(logging.INFO)
mngr_conf = {'root_log_directory': '/tmp'}
lm = bsc.StreamCaptureManager(mngr_conf, [])
bad_func_name = 'this function name doesnt exist'
kwargs = {'pre_write_transforms': [bad_func_name]}
lm.add_logger('testlog', ['', 9876], 'udp', '~/logger_dir', **kwargs)
msg = 'Unable to load data transformation "{}" for handler "{}"'.format(
bad_func_name, 'testlog')
log_mock.assert_called_with(msg)
logging.getLogger('ait').setLevel(logging.CRITICAL)
def test_bad_type_transform_load(self, socket_mock, mkdirs_mock, open_mock,
log_mock):
logging.getLogger("ait").setLevel(logging.INFO)
mngr_conf = {"root_log_directory": "/tmp"}
lm = bsc.StreamCaptureManager(mngr_conf, [])
bad_func_name = ("foobarbaz", )
kwargs = {"pre_write_transforms": [bad_func_name]}
lm.add_logger("testlog", ["", 9876], "udp", "~/logger_dir", **kwargs)
msg = 'Unable to determine how to load data transform "{}"'.format(
bad_func_name)
log_mock.assert_called_with(msg)
logging.getLogger("ait").setLevel(logging.CRITICAL)
def test_bad_builtin_transform_load(self, socket_mock, mkdirs_mock,
open_mock, log_mock):
logging.getLogger("ait").setLevel(logging.INFO)
mngr_conf = {"root_log_directory": "/tmp"}
lm = bsc.StreamCaptureManager(mngr_conf, [])
bad_func_name = "this function name doesnt exist"
kwargs = {"pre_write_transforms": [bad_func_name]}
lm.add_logger("testlog", ["", 9876], "udp", "~/logger_dir", **kwargs)
msg = 'Unable to load data transformation "{}" for handler "{}"'.format(
bad_func_name, "testlog")
log_mock.assert_called_with(msg)
logging.getLogger("ait").setLevel(logging.CRITICAL)
def test_pre_write_transform_load(self, socket_mock, mkdirs_mock,
pcap_open_mock):
mngr_conf = {'root_log_directory': '/tmp'}
lm = bsc.StreamCaptureManager(mngr_conf, [])
kwargs = {'pre_write_transforms': ['identity_transform', lambda x: 1]}
lm.add_logger('testlog', ['', 9876], 'udp', '~/logger_dir', **kwargs)
stream_capturer = lm._stream_capturers["['', 9876]"][0]
handler = stream_capturer.capture_handlers[0]
assert 'pre_write_transforms' in handler
assert len(handler['pre_write_transforms']) == 2
for t in handler['pre_write_transforms']:
assert hasattr(t, '__call__')
assert 'identity_transform' == handler['pre_write_transforms'][
0].__name__
assert 1 == handler['pre_write_transforms'][1]('bogus input')
def test_pre_write_transform_load(self, socket_mock, mkdirs_mock,
pcap_open_mock):
mngr_conf = {"root_log_directory": "/tmp"}
lm = bsc.StreamCaptureManager(mngr_conf, [])
kwargs = {"pre_write_transforms": ["identity_transform", lambda x: 1]}
lm.add_logger("testlog", ["", 9876], "udp", "~/logger_dir", **kwargs)
stream_capturer = lm._stream_capturers["['', 9876]"][0]
handler = stream_capturer.capture_handlers[0]
assert "pre_write_transforms" in handler
assert len(handler["pre_write_transforms"]) == 2
for t in handler["pre_write_transforms"]:
assert hasattr(t, "__call__")
assert "identity_transform" == handler["pre_write_transforms"][
0].__name__
assert 1 == handler["pre_write_transforms"][1]("bogus input")
def test_add_logger(self, socket_log_mock, mkdirs_mock):
mngr_conf = {"root_log_directory": "/totally/a/../fake/dir/../name"}
# We'll use this to make sure directory paths are cleaned
cleaned_dir_path = os.path.normpath(mngr_conf["root_log_directory"])
lm = bsc.StreamCaptureManager(mngr_conf, [])
lm.add_logger("foo", ["", 9000], "udp", "/tmp")
assert len(lm._stream_capturers.keys()) == 1
assert "['', 9000]" in lm._stream_capturers
# Default root_log_directory usage and normalization check
lm.add_logger("baz", ["", 8500], "udp")
socket_log_mock.assert_called_with(
{
"log_dir": cleaned_dir_path,
"name": "baz",
"rotate_log": True,
"pre_write_transforms": [],
},
["", 8500],
"udp",
)
assert lm._pool.free_count() == 48
# Check to make sure that home directory expansion is being done
socket_log_mock.reset_mock()
lm.add_logger("testlog", ["", 1234], "udp", "~/logger_dir")
expanded_user_path = os.path.expanduser("~/logger_dir")
socket_log_mock.assert_called_with(
{
"log_dir": expanded_user_path,
"name": "testlog",
"rotate_log": True,
"pre_write_transforms": [],
},
["", 1234],
"udp",
)
def test_remove_logger(self, socket_log_mock):
lm = bsc.StreamCaptureManager(None, [])
lm.add_logger('foo', ['', 9000], 'udp', '/tmp')
lm.stop_capture_handler('foo')
assert mock.call().remove_handler('foo') in socket_log_mock.mock_calls
def test_remove_logger(self, socket_log_mock):
lm = bsc.StreamCaptureManager(None, [])
lm.add_logger("foo", ["", 9000], "udp", "/tmp")
lm.stop_capture_handler("foo")
assert mock.call().remove_handler("foo") in socket_log_mock.mock_calls
