#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import argparse
import logging
import logging.config
import os.path

from config.constant import Values
import components.data
import components.train
import components.predict
import components.evaluate
import components.group_filter
import strfilters.language_filter
import strfilters.tokenizer
from semantic.taint import GroupLabel
import semantic.taint
import argparse
import logging
import logging.config
import os.path
import visitor.loader
import visitor.saver
import sqlite3


index =0
def split_train_test_malware_b(train_all,test_all):
	logger = logging.getLogger(__name__)

	if train_all is None:
		train_all = accessor.loader.JsonLoader.load(Values.all_train_labeled_tainted_groups_multiIMG_path)
		test_all = accessor.loader.JsonLoader.load(Values.all_test_labeled_tainted_groups_multiIMG_path)

	global index
	logger.info("++++++++++++++"+str(index)+"+++++++++++++++")

	db_connection = sqlite3.connect('/home/Research/results/datas_old/screen_api_flow_db_filter_3.0.db')
	cursor = db_connection.execute("select * from screenapiflow where label = ?",('b',))
	fdroid_benign = list()
	for one in cursor:
		if 'malware/' in one[1]: 
			fdroid_benign.append(one)
	print(len(fdroid_benign))
	app2benigns = dict()
	for benign in fdroid_benign:
		if benign[1] in app2benigns:
			ids = app2benigns[ benign[1]]
		else:
			ids = list()
		ids.append(benign[0])
		app2benigns[benign[1]] = ids

	train_b = list()
	test_b = list()
	temp = list()
	for key in app2benigns.keys():
		ids = app2benigns[key]
		if len(ids) == 1:
			temp.extend(ids)
		else:
			if len(ids) == 2:
				train, test = modules.data.split_set(groups=ids, test_percentage=0.5, shuffle=True)
			else:
				train, test = modules.data.split_set(groups=ids, test_percentage=0.4, shuffle=True)
			train_b.extend(train)
			test_b.extend(test)

	train, test = modules.data.split_set(groups=temp, test_percentage=0.45, shuffle=True)
	train_b.extend(train)
	test_b.extend(test)
	print(len(train_b))
	print(len(test_b))
	# train_all = accessor.loader.JsonLoader.load(Values.all_train_labeled_tainted_groups_multiIMG_path)
	# test_all = accessor.loader.JsonLoader.load(Values.all_test_labeled_tainted_groups_multiIMG_path)
	exit()
	id2group = dict()
	for group in train_all[:]:
		app_descr = group['app_descr']
		label = group['label']
		if len(app_descr['raw'])==0 and "BENIGN" in label:
			id2group[str(group['id'])] = group
			train_all.remove(group)



	for group in test_all[:]:
		app_descr = group['app_descr']
		label = group['label']
		if len(app_descr['raw'])==0 and "BENIGN" in label:
			id2group[str(group['id'])] = group
			test_all.remove(group)



	for tid in train_b:
		tid = str(tid)
		if tid in id2group:
			train_all.append(id2group[tid])
		else:
			print('id:'+tid)
	for tid in test_b:
		tid = str(tid)
		if tid in id2group:
			test_all.append(id2group[tid])
		else:
			print('id:'+tid)

	logger.info("++++++++++++++"+str(len(train_all))+"+++++++++++++++")
	logger.info("++++++++++++++"+str(len(test_all))+"+++++++++++++++")

	accessor.saver.JsonSaver.save(train_all,Values.tmp_multiIMG_train_path+str(index))
	accessor.saver.JsonSaver.save(test_all,Values.tmp_multiIMG_test_path+str(index))

	db_connection.close()

	train = [entity.taint.Group().parse_json(item) for item in train_all]
	test = [entity.taint.Group().parse_json(item) for item in test_all]

	from model.test_best_model import FeatureModel
	best_model = FeatureModel()
	best_model.test_best_model_whole(train,test)

def split_train_test_benign_b():
	# data = accessor.loader.JsonLoader.load(Values.all_labeled_tainted_groups_multiIMG_path)
	# train, test = modules.data.split_set(groups=data, test_percentage=0.5,
	# 	                                       shuffle=True)
	# train = accessor.saver.JsonSaver.save(train,Values.all_train_labeled_tainted_groups_multiIMG_path)
	# test = accessor.saver.JsonSaver.save(test,Values.all_test_labeled_tainted_groups_multiIMG_path)

	db_connection = sqlite3.connect('/home/Research/results/datas_old/screen_api_flow_db_filter_3.0.db')
	cursor = db_connection.execute("select * from screenapiflow where label = ?",('b',))
	fdroid_benign = list()
	for one in cursor:
		if 'fdroid/' in one[1]: 
			fdroid_benign.append(one)
	print(len(fdroid_benign))
	app2benigns = dict()
	for benign in fdroid_benign:
		if benign[1] in app2benigns:
			ids = app2benigns[ benign[1]]
		else:
			ids = list()
		ids.append(benign[0])
		app2benigns[benign[1]] = ids

	train_b = list()
	test_b = list()
	temp = list()
	for key in app2benigns.keys():
		ids = app2benigns[key]
		if len(ids) == 1:
			temp.extend(ids)
		else:
			train, test = modules.data.split_set(groups=ids, test_percentage=0.5, shuffle=True)
			train_b.extend(train)
			test_b.extend(test)

	train, test = modules.data.split_set(groups=temp, test_percentage=0.3, shuffle=True)
	train_b.extend(train)
	test_b.extend(test)
	print(len(train_b))
	print(len(test_b))
	train = accessor.loader.JsonLoader.load(Values.all_train_labeled_tainted_groups_multiIMG_path)
	test = accessor.loader.JsonLoader.load(Values.all_test_labeled_tainted_groups_multiIMG_path)

	id2group = dict()
	for group in train[:]:
		app_descr = group['app_descr']
		label = group['label']
		if len(app_descr['raw'])>0 and "BENIGN" in label:
			id2group[str(group['id'])] = group
			train.remove(group)



	for group in test[:]:
		app_descr = group['app_descr']
		label = group['label']
		if len(app_descr['raw'])>0 and "BENIGN" in label:
			id2group[str(group['id'])] = group
			test.remove(group)



	for tid in train_b:
		tid = str(tid)
		if tid in id2group:
			train.append(id2group[tid])
		else:
			print('id:'+tid)
	for tid in test_b:
		tid = str(tid)
		if tid in id2group:
			test.append(id2group[tid])
		else:
			print('id:'+tid)
	print(len(train))
	print(len(test))

	db_connection.close()

	split_train_test_malware_b(train,test)
#	exit()

#	train = [entity.taint.Group().parse_json(item) for item in train]
#	test = [entity.taint.Group().parse_json(item) for item in test]

#	from model.test_best_model import FeatureModel
#	best_model = FeatureModel()
#	best_model.test_best_model_whole(train,test)
def split_train_test_data():
	split_train_test_benign_b()


def read_ui_state_pairs():
	db_connection = sqlite3.connect('/home/Research/tools/results/datas_old/screen_api_flow_db_filter_3.0.db')
	cursor = db_connection.execute("select * from screenapiflow")
	senapi = dict()
	screenimagepath = set()
	apps = set()
	index = 1
	for one in cursor:
		apps.add(one[1])
		if one[3] not in senapi:
			senapi[one[3]] = 1
		else:
			senapi[one[3]] += 1
		screenimagepath.add(one[9])
	print(len(senapi))
	for api in senapi:
		print(api,senapi[api])
	print(len(screenimagepath))
	print(len(apps))

	db_connection.close()

def compute_categories():
	data = accessor.loader.JsonLoader.load(Values.all_filter_train_labeled_tainted_groups_multiIMG_path)

	monitor_apis = accessor.loader.JsonLoader.load(Values.monitor_recorded_apis_path)

	monitor_recorded_apis = dict()
	for key in monitor_apis:
		values = monitor_apis[key]
		monitor_recorded_apis[key] = set(line.strip() for line in values.split(', '))

	# data = accessor.loader.JsonLoader.load(Values.all_filter_test_labeled_tainted_groups_multiIMG_path)
	print(len(data))
	for one_type in monitor_recorded_apis.keys():
		print("********"+one_type+"********")
		one_type_groups = list()
		for group in data:
			flow = group['flow']
			raw = flow['raw']
			raw = raw.split('=>')[1]
			mtd,clss = raw.split('@')
			api = clss+"."+mtd
			if api in monitor_recorded_apis[one_type]:
				one_type_groups.append(group)
		
		# print(len(one_type_groups))
		one_type_groups = [entity.taint.Group().parse_json(item) for item in one_type_groups]
		benigns = set()
		malwares = set()
		for group in one_type_groups:
			if group.label.name == GroupLabel.BENIGN.name:
				benigns.add(group)
			if group.label.name == GroupLabel.MALICIOUS.name:
				malwares.add(group)
		print(len(one_type_groups),len(benigns),len(malwares))
		# exit()

def split_exist_train_test_data():
	train_path = Values.all_filter_train_labeled_tainted_groups_multiIMG_path.replace('apicogdata','apidata')
	exist_train = visitor.loader.JsonLoader.load(train_path)
	test_path = Values.all_filter_test_labeled_tainted_groups_multiIMG_path.replace('apicogdata','apidata')
	exist_test = visitor.loader.JsonLoader.load(test_path)

	train_ids = set()
	for group in exist_train[:]:
		train_ids.add(group['id'])
	
	# train, test = components.data.split_set(groups=data, test_percentage=0.5,
	# 	                                       shuffle=True)
	data = visitor.loader.JsonLoader.load(Values.all_filtered_labeled_tainted_groups_path)

	train = list()
	test = list()
	for group in data:
		if group['id'] in train_ids:
			train.append(group)
		else:
			test.append(group)
	print(len(train),len(test))
	visitor.saver.JsonSaver.save(train,Values.all_filter_train_labeled_tainted_groups_multiIMG_path)
	visitor.saver.JsonSaver.save(test,Values.all_filter_test_labeled_tainted_groups_multiIMG_path)

if __name__ == '__main__':
	# logging.config.dictConfig(Values.log_init_config)
	# logger = logging.getLogger(__name__)
	# for indx in range(10):
	# 	index = indx
	# 	split_train_test_data()

	# split_train_test_malware_b(None,None)

	# train = accessor.loader.JsonLoader.load(Values.tmp_multiIMG_train_path+str(3))
	# test = accessor.loader.JsonLoader.load(Values.tmp_multiIMG_test_path+str(3))

	# train = [entity.taint.Group().parse_json(item) for item in train]
	# test = [entity.taint.Group().parse_json(item) for item in test]

	# from model.test_best_model import FeatureModel
	# best_model = FeatureModel()
	# best_model.test_best_model_whole(train,test)


	
	# from gensim.models import KeyedVectors
	# model = KeyedVectors.load(Values.word2vec_model_path)
	# try:
	# 	similar_res = [x[0] for x in model.wv.most_similar('open',topn=50)]
	# except (KeyError):
	# 	similar_res = set()
	# print(similar_res)
	split_exist_train_test_data()