diff --git a/hp-labrea/labrea-idea.py b/hp-labrea/labrea-idea.py index 2f3b170b29e661460e0ffea10ef8f3b8e7afcdd7..0afb6ed48c6191a7496c08010126d4a3be47c4e5 100755 --- a/hp-labrea/labrea-idea.py +++ b/hp-labrea/labrea-idea.py @@ -7,7 +7,6 @@ import re import time import optparse import signal -import contextlib import uuid import json import socket @@ -28,6 +27,17 @@ try: except ImportError: from ordereddict import OrderedDict +try: + from contextlib import nested +except ImportError: + from contextlib import ExitStack, contextmanager + + @contextmanager + def nested(*contexts): + with ExitStack() as stack: + for ctx in contexts: + stack.enter_context(ctx) + yield contexts class WindowContextMgr(object): @@ -46,7 +56,7 @@ class WindowContextMgr(object): def expire_queue(self, queue, window): aggr_events = [] ctx_to_del = [] - for ctx, timestamp in queue.iteritems(): + for ctx, timestamp in queue.items(): if timestamp >= self.update_timestamp - window: break closed = self.ctx_close(self.contexts[ctx]) @@ -304,7 +314,7 @@ class IdeaGen(object): def gen_idea(self, src, src_ports, targets, detect_time, event_time, cease_time, count, template): tmpl = self.template[template] isource = { - "IP6" if ":" in src else "IP4": [src], + "IP4": [src], "Proto": tmpl["proto"] } if "source_type" in tmpl: @@ -318,12 +328,8 @@ class IdeaGen(object): itargets = [] for ports, tgt in folded_tgt.items(): itarget = {"Proto": tmpl["proto"]} - tgts4 = [ip for ip in tgt if ":" not in ip] - tgts6 = [ip for ip in tgt if ":" in ip] - if tgts4: - itarget["IP4"] = tgts4 - if tgts6: - itarget["IP6"] = tgts6 + if tgt: + itarget["IP4"] = tgt if ports: itarget["Port"] = [int(port) for port in ports] itargets.append(itarget) @@ -661,7 +667,7 @@ def main(): for context in [PingContextMgr, ConnectContextMgr, InboundContextMgr] ] - with contextlib.nested(*files): + with nested(*files): for line_set in izip(*files): for line in line_set: if not line: