diff --git a/prometheus_ss_exporter/selection.py b/prometheus_ss_exporter/selection.py index 19dd714..4c5e16d 100644 --- a/prometheus_ss_exporter/selection.py +++ b/prometheus_ss_exporter/selection.py @@ -6,11 +6,14 @@ class Selector: class Discerner: def ports(self, flow, portranges): + if not portranges: + return True + for p_range in portranges: - if ((flow['dst_port'] < p_range['lower']) or - (flow['dst_port'] > p_range['upper'])): - return False - return True + if ((flow['dst_port'] >= p_range['lower']) and + (flow['dst_port'] <= p_range['upper'])): + return True + return False def peers(self, flow, hosts=[], addresses=[], networks=[]): @@ -41,10 +44,11 @@ def process(self, flow, pids=[], cmds=[]): flow_pids = [] flow_cmds = [] - for usr, pid_ctxt in flow['usr_ctxt'].items(): - for pid, cmd_ctxt in pid_ctxt.items(): - flow_pids.append(pid) - flow_cmds.append(cmd_ctxt['full_cmd']) + if flow.get('usr_ctxt'): + for usr, pid_ctxt in flow['usr_ctxt'].items(): + for pid, cmd_ctxt in pid_ctxt.items(): + flow_pids.append(pid) + flow_cmds.append(cmd_ctxt['full_cmd']) for pid in pids: if pid in flow_pids: @@ -71,16 +75,15 @@ def arbitrate(self, flow): return self._core(flow) def _arbitrate(self, flow): - conditions = [self.discern.ports(flow, self.cnfg['peering']['portranges']), - self.discern.peers(flow, - hosts=self.cnfg['peering']['nodes']['hosts'], - addresses=self.cnfg['peering']['nodes']['addresses'], - networks=self.cnfg['peering']['nodes']['networks']), - self.discern.process(flow, - pids=self.cnfg['process']['pids'], - cmds=self.cnfg['process']['cmds']) + conditions = [ self.discern.ports(flow, self.cnfg.get('peering').get('portranges')) if self.cnfg.get('peering') else True, + self.discern.peers(flow, hosts=self.cnfg.get('peering').get('hosts'), + addresses=self.cnfg.get('peering').get('addresses'), + networks=self.cnfg.get('peering').get('networks')) if self.cnfg.get('peering') else True, + self.discern.process(flow, pids=self.cnfg.get('process').get('pids'), + cmds=self.cnfg.get('process').get('cmds')) if self.cnfg.get('process') else True ] - if it.dropwhile(lambda _: _, conditions): - return True + # if one condition false, we decline sample + if list(it.filterfalse(lambda _: _, conditions)): + return False - return False + return True diff --git a/test/selection_test.py b/test/selection_test.py index 69e16a0..a5aeded 100644 --- a/test/selection_test.py +++ b/test/selection_test.py @@ -8,6 +8,7 @@ class SelectorTesting(unittest.TestCase): @classmethod def setUpClass(cls): cls.discerner = selection.Selector.Discerner() + cls.selector = selection.Selector def test_peers_decline(self): selector_addr = ["10.0.1.10"] @@ -94,3 +95,35 @@ def test_process_accept_test(self): outcome = SelectorTesting.discerner.process(flow, cmds=selector_cmds) self.assertTrue(outcome) + + def test_combined_accept_test(self): + selection_config = { + 'selection': { + 'peering': { + 'networks': [ '192.168.0.0/16' ], + 'portranges': [{ 'lower': 1000, 'upper': 2000 }] + } + } + } + flow = {'dst': '192.168.92.41', 'dst_port': 1500} + + selector = SelectorTesting.selector(selection_config) + outcome = selector.arbitrate(flow) + + self.assertTrue(outcome) + + def test_combined_decline_test(self): + selection_config = { + 'selection': { + 'peering': { + 'networks': [ '192.168.0.0/16' ], + 'portranges': [{ 'lower': 1000, 'upper': 2000 }] + } + } + } + flow = {'dst': '192.168.92.41', 'dst_port': 900} + + selector = SelectorTesting.selector(selection_config) + outcome = selector.arbitrate(flow) + + self.assertFalse(outcome)