"""
2015 - Mike Schladt
Simple helper script to add vectors to mfv database

VECTOR TABLE DEF :

CREATE TABLE vectors(
    md5 CHAR(32) NOT NULL,
	machine TEXT NOT NULL,
	task_id INT UNSIGNED NOT NULL,
	type TEXT NOT NULL, 
	value INT NOT NULL
);
    
VECTORS FILE FORMAT : default output of create_feature_vectors.py 
"""

import MySQLdb 
import argparse
import sys

db_host = CONST.DB_HOST
db_user = CONST.DB_USER
db_pass = CONST.DB_PASS
db_name = CONST.DB_NAME

connection = MySQLdb.connect(db_host, db_user, db_pass, db_name)
cursor = connection.cursor()

parser = argparse.ArgumentParser(description='Simple helper script to vectors to malware feature vector database')
parser.add_argument('vector_file', help='Filepath of vector files', nargs='+')
args = parser.parse_args()

#submit samples to cuckoo sanbox
for file in args.vector_file :
    lines = open(file).readlines()
    
    #get keywords from first line
    keys = lines.pop(0).strip()
    keys = keys.split(',')
    keys = keys[3:]
    
    #create vector db entries for each task / line
    for line in lines : 
        vector = line.split(',')
        md5 = vector.pop(0).strip()
        machine_name = vector.pop(0).strip()
        task_id = vector.pop(0).strip()
        """
        DEBUG
        print "MD5 : {0}".format(md5)
        print "Machine Name : {0}".format(machine_name)
        print "Task ID : {0}".format(task_id)
        """
        i = 0
        for item in vector : 
            vector_type = keys[i].strip()
            vector_value = item.strip()
            i += 1
            #DEBUG print "Vector Type : {0} , Vector Value : {1}".format(vector_type, vector_item)
            
            #add vector entry to db
            cmd = "INSERT INTO vectors (md5, machine, task_id, type, value) VALUES (%s, %s, %s, %s, %s)"
            data = [md5, machine_name, task_id, vector_type, vector_value]
            try : 
                cursor.execute(cmd, data)
            except Exception as e:
                if "Duplicate entry" in str(e) :
                    continue
                else : 
                    raise    
                    
        print "Added vectors for task : {0} , MD5 {1}".format(task_id, md5)
               
#commit and close database connection 
connection.commit()
connection.close()