#!/usr/bin/env python2

__author__ = "Andrea Fioraldi"
__mail__ = "fioraldi.1692419@studenti.uniroma1.it"

import urlparse
import random
import json
import os

#to_consider = ["feature", "permission", "intent", "api_call", "real_permission", "activity", "url", "call"]
#to_consider = ["feature", "intent", "api_call", "real_permission", "call"]
#to_consider = ["api_call", "feature", "intent"]
#to_consider = ["real_permission", "api_call"]
to_consider = ["feature", "intent", "api_call", "real_permission"]
#to_consider = ["permission", "url", "service", "feature"]

random.seed()

print (" >> reading the dataset...")

with open("./drebin/sha256_family.csv") as f:
    malicious = set(map(lambda x: x.strip().split(",")[0], filter(lambda x: len(x) > 0, f.readlines())[1:]))

dataset = set(os.listdir("./drebin/feature_vectors/"))
goods = dataset.difference(malicious)

print (" >> sampling...")

validation_malicious = random.sample(malicious, len(malicious) / 3)
#validation_goods = random.sample(goods, len(malicious) / 3)
validation_goods = random.sample(goods, len(goods) / 3)

training_malicious = malicious.difference(validation_malicious)
training_goods = goods.difference(validation_goods)

features_map = {}

print (" >> processing...")

validation = ""
for e in validation_malicious:
    d = "malware,"
    
    with open("./drebin/feature_vectors/" + e) as f:
        for x in filter(lambda x: x[0] in to_consider, map(lambda x: x.strip().split("::", 1), f.readlines())):
            if x[0] == "url":
                try:
                    x[1] = urlparse.urlparse(x[1]).netloc
                except:
                    pass
            feat = x[0] + "::" + x[1].replace(",", "")
            #features_map[feat] = features_map.get(feat, 0) + 1
            d += feat + ","
        
    validation += d[:-1] + "\n"

for e in validation_goods:
    d = "notmalware,"
    
    with open("./drebin/feature_vectors/" + e) as f:
        for x in filter(lambda x: x[0] in to_consider, map(lambda x: x.strip().split("::", 1), f.readlines())):
            if x[0] == "url":
                try:
                    x[1] = urlparse.urlparse(x[1]).netloc
                except:
                    pass
            feat = x[0] + "::" + x[1].replace(",", "")
            #features_map[feat] = features_map.get(feat, 0) + 1
            d += feat + ","
    
    validation += d[:-1] + "\n"

training = []
for e in training_malicious:
    d = ["malware"]
    
    with open("./drebin/feature_vectors/" + e) as f:
        for x in filter(lambda x: x[0] in to_consider, map(lambda x: x.strip().split("::", 1), f.readlines())):
            if x[0] == "url":
                try:
                    x[1] = urlparse.urlparse(x[1]).netloc
                except:
                    pass
            feat = x[0] + "::" + x[1].replace(",", "")
            features_map[feat] = features_map.get(feat, 0) + 1
            d.append(feat)
        
    training.append(d)

for e in training_goods:
    d = ["notmalware"]
    
    with open("./drebin/feature_vectors/" + e) as f:
        for x in filter(lambda x: x[0] in to_consider, map(lambda x: x.strip().split("::", 1), f.readlines())):
            if x[0] == "url":
                try:
                    x[1] = urlparse.urlparse(x[1]).netloc
                except:
                    pass
            feat = x[0] + "::" + x[1].replace(",", "")
            features_map[feat] = features_map.get(feat, 0) + 1
            d.append(feat)
        
    training.append(d)


for i, x in enumerate(training):
    training[i] = [x[0]] + filter(lambda x: features_map[x] > 1, x[1:])
training = filter(lambda x: len(x) > 1, training)
training = map(",".join, training)

print " >> saving..."

with open("validation.csv", "w") as f:
    f.write(validation)

with open("training.csv", "w") as f:
    f.write("\n".join(training))

with open("feature_classes.txt", "w") as f:
    f.write(" ".join(to_consider))

