'''
Author: Ashish Katlam
Description: Fetches data from database and create a csv file
'''
from pymongo import MongoClient
from constants import DB_NAME
from core2 import create_vector_single


client = MongoClient()
db = client[DB_NAME]

permissions = []
apis = []

# Get unique permissions
with open('unique_permissions.txt','r') as fp1:
	for line in fp1:
		permissions.append(line.strip().replace(",","").replace("'",""))

# Get apis
with open('api.txt','r') as fp2:
	for line in fp2:
		apis.append(line.strip().replace(",","").replace("'",""))

features = permissions + apis
features.append('com.metasploit.stage.PayloadTrustManager')
features.append('entropy_rate')
features.append('native')
features.append('db')
features.append('class')

with open('data.csv','w+') as op:
	header = ""
	for f in features:
		header+= f.strip().replace('"','')+','
	header = header[:-1]
	op.write(header+'\n')


	for apk in db.apk.find():
		feature_vector = create_vector_single(apk)
		str_to_write = ""
		for i,feature in enumerate(feature_vector):
			if i < len(feature_vector)-1:
				str_to_write+=str(feature)+','
			else:
				class_label = 1 if apk['data_type'] == 'malware' else 0
				str_to_write+=str(feature)+','+str(class_label)
		op.write(str_to_write+'\n')




