MalwareClassifier/bin/import.py

import argparse
import sys
import redis

argParser = argparse.ArgumentParser(description='Malware classifier')
argParser.add_argument('-f', action='append', help='Filename')
args = argParser.parse_args()

r = redis.StrictRedis(host='localhost', port=6379, db=0)

if args.f is not None:
    md5 = args.f[0].split(".")[0]
    r.sadd('processed', md5)
    lnumber = 0
    fields = None
    for line in sys.stdin:
        if lnumber == 0:
            fields = line.rstrip().split(",")
            for field in fields:
                r.sadd('type', field)
        else:
            elements = line.rstrip().split(",")
            i = 0
            for element in elements:
                try:
                    r.sadd('e:'+fields[i], element)
                except IndexError:
                    print("Empty fields")
                i = i+1

        lnumber = lnumber + 1
else:
    argParser.print_help()