diff --git a/flowmon-ads/warden3_flowmon_ads_filer.py b/flowmon-ads/warden3_flowmon_ads_filer.py index 1ab1c2a686c637fbe52a0041ed09734f907bd51b..dc449be4fd53f28769aa0f0469e243447611f7d3 100644 --- a/flowmon-ads/warden3_flowmon_ads_filer.py +++ b/flowmon-ads/warden3_flowmon_ads_filer.py @@ -6,6 +6,7 @@ import os import sys +import getopt sys.path.append('/data/warden/libs') from warden_client import read_cfg, format_time @@ -15,6 +16,42 @@ import csv from time import strptime, mktime from uuid import uuid4 +# Command line options handling +# Had to use getopt for 2.6 compatibility. Meh. :( + +opt_dict = { + "help": False, + "test": False, + "origdata": False, + "errlog": "/data/warden/var/flowmon-ads-filer_lastrun.log", + "out": "/data/warden/var/feeds-out", + "target": "NONE" +} + +getopt_format = [k + (v and "=" or "") for (k, v) in opt_dict.items()] + +def help(s=None, exitcode=0): + if s is not None: + print("Error: %s" % s) + print("Usage: %s [%s]" % (sys.argv[0], "] [".join(["--" + v for v in getopt_format]))) + sys.exit(exitcode) + +def get_opts(): + try: + rawopts, args = getopt.getopt(sys.argv[1:], "", getopt_format) + except getopt.GetoptError as err: + help(err, exitcode=2) + if args: + help("unknown arg %s" % ", ".join(args), exitcode=2) + rawopts = dict(rawopts) + if "--help" in rawopts: + help() + opts = {} + for k, v in opt_dict.items(): + opts[k] = v and rawopts.get("--" + k, v) or ("--" + k) in rawopts + return opts + + # Conversion/validation routines def isotime(t): if not t: @@ -128,7 +165,7 @@ ads_types = { def xlat_ads_type(s): if s not in ads_types.keys(): return [] - return ads_types[s] + return ads_types[s][:] def xlat_ads_proto(s): @@ -141,7 +178,7 @@ def xlat_ads_proto(s): return proto -def gen_idea_from_ads(ads, orig_data, anonymised_target): +def gen_idea_from_ads(ads, orig_data, anonymised_target, add_test): # Mandatory ts = ads["Timestamp"] or time.localtime() @@ -152,6 +189,9 @@ def gen_idea_from_ads(ads, orig_data, anonymised_target): "DetectTime": format_time(*ts[0:6]) } + if add_test: + event["Category"].append("Test") + # Optional if ads["ID"]: event["AltNames"] = ["ADS-%i" % ads["ID"]] @@ -185,7 +225,7 @@ def gen_idea_from_ads(ads, orig_data, anonymised_target): if ads["Protocol"]: target["Proto"] = [xlat_ads_proto(p) for p in ads["Protocol"]] - if anonymised_target: + if anonymised_target != "NONE": tgtips = [anonymised_target] else: tgtips = ads["Targets"] @@ -215,9 +255,13 @@ def gen_idea_from_ads(ads, orig_data, anonymised_target): def main(): - anonymised_target = None - add_orig_data = True - sdir = SafeDir("/data/warden/var/feeds-out") + opts = get_opts() + try: + errlog = open(opts["errlog"], "w") + sys.stderr = errlog + except IOError: + print("Warning: error log %s unavailable (wrong directory or permissions?)" % opts["errlog"]) + out = SafeDir(opts["out"]) ads_fields = [it[0] for it in sorted(ads_fieldnames.items(), key=lambda it: it[1]["order"])] for row in csv.reader(sys.stdin, dialect="excel-tab"): @@ -226,10 +270,10 @@ def main(): tr_row = {} for k, val in zip(ads_fields, row): tr_row[k] = xlat_ads_field(k, val) - if not add_orig_data: + if not opts["origdata"]: row = None - event = gen_idea_from_ads(tr_row, row, anonymised_target) - nf = sdir.newfile() + event = gen_idea_from_ads(tr_row, row, opts["target"], opts["test"]) + nf = out.newfile() try: data = json.dumps(event) nf.f.write(data.encode("utf-8")) @@ -238,10 +282,8 @@ def main(): sys.stderr.write("Error source line: %s\n" % row) sys.stderr.write("Error event data: %s\n" % str(event)) nf.f.close() - nf.moveto(sdir.incoming) + nf.moveto(out.incoming) if __name__ == "__main__": - if os.path.exists('/data/warden/var'): - sys.stderr = open('/data/warden/var/flowmon-ads-filer_lastrun.log', 'w') main()