Added possibility to add "Test" category; implemented Py 2.6 compatible cmd line options

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()