from AttackerCalc import AttackerCalc
from FeaturesCalc import FeaturesCalc
from PacketFilter import PacketFilter
from MTADownloader import MTADownloader
from NeuralNetwork import NeuralNetwork
from LegitimateDownloader import LegitimateDownloader
from Agent import Agent
from CSV import CSV
from scapy.all import *
import os, signal, sys, glob, time
import gym
import gym_DRLMal

class StartTrainingHandler():

    def __init__(self, nb_malware_training=1, nb_legitimate_training=1, window_size=10, single_csv=False):
        assert window_size > 0, "Inserire un valore valido per window size"
        assert nb_malware_training > 0, "Inserire un valore valido per nb_malware_training"
        assert nb_legitimate_training > 0, "Inserire un valore valido per nb_legitimate_training"
        assert (single_csv is True) or (single_csv is False), "Inserire un valore valido per single_csv"
        self.single_csv = single_csv
        self.nb_malware_training = nb_malware_training
        self.nb_legitimate_training = nb_legitimate_training
        self.window_size = window_size
        self.mta = MTADownloader(save_downloaded=False)
        self.lg = LegitimateDownloader()
        self.ip_to_ignore = ["127.0.0.1"]
        self.trained_pcaps = "trained_pcaps.txt"
        self.featuresCalc = FeaturesCalc(flow_type="malware", min_window_size=self.window_size)
        filter_1 = PacketFilter(ip_whitelist_filter=[], ip_blacklist_filter=self.ip_to_ignore, TCP=True)
        filter_2 = PacketFilter(ip_whitelist_filter=[], ip_blacklist_filter=self.ip_to_ignore, UDP=True)
        filter_3 = PacketFilter(ip_whitelist_filter=[], ip_blacklist_filter=self.ip_to_ignore, ICMP=True)
        self.filters = [filter_1, filter_2, filter_3]

        # Env, Agent and NN definition
        ENV_NAME = 'DRLMal-v0'
        self.env = gym.make(ENV_NAME)
        nn = NeuralNetwork(input_shape=(1,) + self.env.observation_space.shape, output_dim=self.env.get_action_space().n)
        model = nn.create_default_mlp_network()
        self.agent = Agent(env=self.env, model=model, num_actions=self.env.get_action_space().n, batch_size=2)
        self.agent.create_default_dqn_agent(nb_steps_warmup=5)

    def start_training(self):

        def check_if_pcap_is_in_trained(pcap):
            try:
                trained_pcaps_file = open(self.trained_pcaps, "r")
                lines = trained_pcaps_file.readlines()
                trovato = False
                pcap = pcap.replace("\n", "")
                for line in lines:
                    line = line.replace("\n", "")
                    if (line == pcap):
                        trovato = True
                        break
                trained_pcaps_file.close()
                return trovato
            except:
                return False

        def add_pcap_in_trained(pcap):
            trained_pcaps_file = open(self.trained_pcaps, "a+")
            trained_pcaps_file.write(pcap + "\n")
            trained_pcaps_file.close()

        def create_malware_features_csv(pcap):
            flow_type = "malware"
            pcap_name = pcap.split("/")
            pcap_name = pcap_name[len(pcap_name) - 1].replace(".pcap", "")
            if (self.featuresCalc.get_flow_type() == flow_type):
                pass
            else:
                self.featuresCalc.set_flow_type(flow_type)
            array_of_pkts = []
            attacker = AttackerCalc(pcap=pcap)
            ip_to_consider = attacker.compute_attacker()
            for filter in self.filters:
                filter.set_ip_whitelist_filter(ip_to_consider)
            pkts = rdpcap(pcap)
            csv = CSV(file_name=pcap_name, folder_name="Features_Malware")
            csv.create_empty_csv()
            csv.add_row(self.featuresCalc.get_features_name())
            filter_res = []
            for pkt in pkts:
                for filter in self.filters:
                    if(filter.check_packet_filter(pkt) is True):
                        filter_res.append(True)
                    else:
                        filter_res.append(False)
                if(True in filter_res):
                    array_of_pkts.append(pkt)
                if (len(array_of_pkts) >= self.featuresCalc.get_min_window_size()):
                    features = self.featuresCalc.compute_features(array_of_pkts)
                    csv.add_row(features)
                    array_of_pkts.clear()
                filter_res.clear()
            csv.close_csv()
            return csv

        def create_legitimate_features_csv(pcap):
            flow_type = "legitimate"
            pcap_name = pcap.split("/")
            pcap_name = pcap_name[len(pcap_name) - 1].replace(".pcap", "")
            if (self.featuresCalc.get_flow_type() == flow_type):
                pass
            else:
                self.featuresCalc.set_flow_type(flow_type)
            array_of_pkts = []
            for filter in self.filters:
                filter.set_ip_whitelist_filter([])
            pkts = rdpcap(pcap)
            csv = CSV(file_name=pcap_name, folder_name="Features_Legitimate")
            csv.create_empty_csv()
            csv.add_row(self.featuresCalc.get_features_name())
            filter_res = []
            for pkt in pkts:
                for filter in self.filters:
                    if(filter.check_packet_filter(pkt) is True):
                        filter_res.append(True)
                    else:
                        filter_res.append(False)
                if(True in filter_res):
                    array_of_pkts.append(pkt)
                if (len(array_of_pkts) >= self.featuresCalc.get_min_window_size()):
                    features = self.featuresCalc.compute_features(array_of_pkts)
                    csv.add_row(features)
                    array_of_pkts.clear()
                filter_res.clear()
            csv.close_csv()
            return csv

        def malware_train(line):
            trained = False
            line = line.replace("\n", "")
            csv_path = ""
            if (check_if_pcap_is_in_trained(line) is False):
                csv = create_malware_features_csv(line)
                csv_path = csv.get_file_path()
                if (self.env.check_history(csv_path) is False):
                    print("\nStart Malware Training\n")
                    self.env.set_csv(csv_path)
                    self.agent.train_agent(steps=self.env.get_rows_number() - 1,
                                           log_interval=self.env.get_rows_number() - 1, verbose=2,
                                           nb_max_episode_steps=self.env.get_rows_number() - 1)
                    self.env.add_csv_to_history(csv_path)
                    add_pcap_in_trained(line)
                    trained = True
                else:
                    print("\n" + csv_path + " gia' utilizzato. Saltato.\n")
                    trained = False
            else:
                print("\n" + line + " gia' utilizzato. Saltato.\n")
                trained = False
            return trained

        def legitimate_train(line):
            trained = False
            line = line.replace("\n", "")
            csv_path = ""
            if(check_if_pcap_is_in_trained(line) is False):
                csv = create_legitimate_features_csv(line)
                csv_path = csv.get_file_path()
                if (self.env.check_history(csv_path) is False):
                    print("\nStart Legitimate Training\n")
                    self.env.set_csv(csv_path)
                    self.agent.train_agent(steps=self.env.get_rows_number() - 1,
                                           log_interval=self.env.get_rows_number() - 1, verbose=2,
                                           nb_max_episode_steps=self.env.get_rows_number() - 1)
                    self.env.add_csv_to_history(csv_path)
                    add_pcap_in_trained(line)
                    trained = True
                else:
                    print("\n" + csv_path + " gia' utilizzato. Saltato.\n")
                    trained = False
            else:
                print("\n" + line + " gia' utilizzato. Saltato.\n")
                trained = False
            return trained

        def compute_features_single_file():
            csv = CSV(file_name="features")
            csv.create_empty_csv()
            csv.add_row(self.featuresCalc.get_features_name())
            mal_lines = glob.glob(self.mta.get_folder_name() + "/" + "*.pcap")
            leg_lines = glob.glob(self.lg.get_folder_name() + "/" + "*.pcap")
            pcaps_in_csv = []
            for line in mal_lines:
                line = line.replace("\n", "")
                if (check_if_pcap_is_in_trained(line) is False):
                    flow_type = "malware"
                    if (self.featuresCalc.get_flow_type() == flow_type):
                        pass
                    else:
                        self.featuresCalc.set_flow_type(flow_type)
                    array_of_pkts = []
                    attacker = AttackerCalc(pcap=line)
                    ip_to_consider = attacker.compute_attacker()
                    for filter in self.filters:
                        filter.set_ip_whitelist_filter(ip_to_consider)
                    pkts = rdpcap(line)
                    filter_res = []
                    for pkt in pkts:
                        for filter in self.filters:
                            if (filter.check_packet_filter(pkt) is True):
                                filter_res.append(True)
                            else:
                                filter_res.append(False)
                        if (True in filter_res):
                            array_of_pkts.append(pkt)
                        if (len(array_of_pkts) >= self.featuresCalc.get_min_window_size()):
                            features = self.featuresCalc.compute_features(array_of_pkts)
                            csv.add_row(features)
                            array_of_pkts.clear()
                        filter_res.clear()
                    add_pcap_in_trained(line)
                    pcaps_in_csv.append(True)
                else:
                    print("\n" + line + " gia' utilizzato. Saltato.\n")
                    pcaps_in_csv.append(False)
            for line in leg_lines:
                line = line.replace("\n", "")
                if (check_if_pcap_is_in_trained(line) is False):
                    flow_type = "legitimate"
                    if (self.featuresCalc.get_flow_type() == flow_type):
                        pass
                    else:
                        self.featuresCalc.set_flow_type(flow_type)
                    array_of_pkts = []
                    for filter in self.filters:
                        filter.set_ip_whitelist_filter([])
                    pkts = rdpcap(line)
                    filter_res = []
                    for pkt in pkts:
                        for filter in self.filters:
                            if (filter.check_packet_filter(pkt) is True):
                                filter_res.append(True)
                            else:
                                filter_res.append(False)
                        if (True in filter_res):
                            array_of_pkts.append(pkt)
                        if (len(array_of_pkts) >= self.featuresCalc.get_min_window_size()):
                            features = self.featuresCalc.compute_features(array_of_pkts)
                            csv.add_row(features)
                            array_of_pkts.clear()
                        filter_res.clear()
                    add_pcap_in_trained(line)
                    pcaps_in_csv.append(True)
                else:
                    print("\n" + line + " gia' utilizzato. Saltato.\n")
                    pcaps_in_csv.append(False)
            if(True in pcaps_in_csv):
                return csv
            else:
                return None

        def single_train():
            trained = False
            csv = compute_features_single_file()
            if(csv is not None):
                csv_path = csv.get_file_path()
                if (self.env.check_history(csv_path) is False):
                    self.env.set_csv(csv_path)
                    self.agent.train_agent(steps=self.env.get_rows_number() - 1,
                                           log_interval=self.env.get_rows_number() - 1, verbose=2,
                                           nb_max_episode_steps=self.env.get_rows_number() - 1)
                    self.env.add_csv_to_history(csv_path)
                    trained = True
                else:
                    print("\n" + csv_path + " gia' utilizzato. Saltato.\n")
                    trained = False
            else:
                trained = False
            return trained

        if(self.single_csv):
            if(single_train()):
                print("\nAddestramento Completato\n")
            else:
                print("\nAddestramento Non Eseguito\n")
        else:
            mal_lines = glob.glob(self.mta.get_folder_name() + "/" + "*.pcap")
            leg_lines = glob.glob(self.lg.get_folder_name() + "/" + "*.pcap")
            lenght_mal_lines = len(mal_lines)
            lenght_leg_lines = len(leg_lines)
            leg_iter = 0
            mal_iter = 0
            mal_counter = 0
            leg_counter = 0
            overflow_mal = False
            overflow_leg = False
            finish_mal = False
            finish_leg = False

            while (True):
                if(mal_counter < lenght_mal_lines):
                    if(mal_iter == self.nb_malware_training):
                        finish_mal = True
                    else:
                        if(malware_train(mal_lines[mal_counter]) is True):
                            mal_iter += 1
                            mal_counter +=1
                        else:
                            mal_counter += 1
                else:
                    overflow_mal = True
                if(leg_counter < lenght_leg_lines):
                    if(leg_iter == self.nb_legitimate_training):
                        finish_leg = True
                    else:
                        if(legitimate_train(leg_lines[leg_counter]) is True):
                            leg_iter += 1
                            leg_counter += 1
                        else:
                            leg_counter += 1
                else:
                    overflow_leg = True

                if(overflow_leg is True and overflow_mal is True):
                    break
                if(finish_mal is True and finish_leg is True):
                    break
                if (overflow_leg is True and finish_mal is True):
                    break
                if (overflow_mal is True and finish_leg is True):
                    break

            print ("\nAddestramento Completato:\n")
            print ("\nMalware Training = " + str(mal_iter) + "\n")
            print ("\nLegitimate Training = " + str(leg_iter) + "\n")
        return
