
import networkx as nx
import matplotlib.pyplot as plt
import matplotlib as mpl
import pandas as pd
import seaborn as sb
import numpy as np
import re
import os

from collections import defaultdict
from Sandy.pp import importProc

class Network:

    def __init__(self, PATH, import_type='mod', compress=False, track_process=False, verbose=False):
        
        assert import_type in {'mod', 'event'}, f"Import type not supported, choose from 'mod' or 'event'; you chose {import_type}"
        
        if import_type == 'mod':
            self.object = importProc.import_mod(PATH, transform=True, compress=compress, track_process=track_process)
        else:
            self.object = importProc.import_event(PATH, transform=True, compress=compress, track_process=track_process, verbose=verbose)
            
        self.PATH = PATH
        self.using_events = False if import_type == 'mod' else True 

    def produce_edge_list(self, source, target, target_node=False):

        self.network = nx.from_pandas_edgelist(self.object,
                                               source=source,
                                               target=target,
                                               edge_attr=True)

    def assign_labels(self, label_name='ProcName'):

        labels = dict(zip(self.network.nodes(), self.object[label_name]))
        self.network = nx.relabel_nodes(self.network, labels, copy=False)

    def assign_attribute(self, attr_name='targetIDX'):

        attr = dict(zip(self.network.nodes(), self.object[attr_name]))
        nx.set_node_attributes(self.network, attr, 'NodeAttr')

        self.attr_colors = create_color_map(self.network, 'NodeAttr')

    def create_target(self, target_name, new_node=False, verbose=False):
        
        if self.using_events:
            from Sandy.pp import importProc
            object = importProc.import_mod(self.PATH, transform=True, compress=False, track_process=False)
        else:
            object = self.object

        # Extract the PID of the parent process
        target = object.loc[object['ProcName'] == target_name, ['PID']]

        target_len = len(target.PID.unique())

        while True:

            target = object.loc[object.PPID.isin(target.PID.unique()) | \
                            object.PID.isin(target.PID.unique()), ['PID']]
            
            # If length stays the same size, the malicious tree must be covered
            if target_len == len(target.PID.unique()):
                
                object = object.assign(targetIDX = object.PID.isin(target.PID).astype(int))
                
                self.target_PID = target.PID.unique()
                self.target_process = object.iloc[target.index].ProcName.unique()

                if self.using_events:
                    self.object = self.object.assign(targetIDX = (self.object).PID.isin(target.PID).astype(int))
                else:
                    self.object = object
                    
                break
          
            else: ## Set new len to old len
                target_len = len(target.PID.unique())
                
        if verbose:
            print(f"Found {((self.object).targetIDX == 1).sum()} target entries")
            print(f"Unique PID for chosen target at: {self.target_PID}")
            print(f"Unique Process Names for chosen target at: {self.target_process}")
            #print(f"Unique Process Names for chosen target at: #['8fc63e8c5de9e771badefaa50cc2f7a0', 'file-2640385_exe']")

        ## Need to work on graphing. As of right now create target does not produce self.network object properly
        if new_node:
            return Network(self.object, self.PATH)
        
        ## Need to make sure it returns node when I need it to
        ## if new_node:
        #net = Network(self.object, self.PATH)
        #net.produce_edge_list("PID", "PPID")
        
    def plot_network(self, plot_attr='targetIDX', save=False, map_feature=False):

        from itertools import count

        fig = plt.figure()
        layout = nx.spring_layout(self.network)

        ## Default to targetIDX if an attribute doesn't exist?
        ## if targetIDX doesnt exist, add it
        #self.assign_attribute('targetIDX')

        ## Retreive create times and set to line widths
        #durations = [int(i[plot_attr]) for i in dict(self.network.edges).values()]
        #min_, max_ = min(durations), max(durations)
        #durations = [(0.1 - 1)*((i - max_)/(min_ - max_)) + 1  for i in durations]

        ## Set node colors to the targetIDX
        groups = set(nx.get_node_attributes(self.network, plot_attr).values())
        mapping = dict(zip(groups, count()))
        nodes = self.network.nodes()
        colors = [mapping[self.network.nodes[n][plot_attr]] for n in nodes]

        nx.draw_networkx_nodes(self.network,pos=layout,node_size=10,
                               node_color=colors,node_list=nodes)
        nx.draw_networkx_edges(self.network,pos=layout,width=durations,style='dashed',
                               edge_color='red',edge_cmap=plt.cm.Blues)
        nx.draw_networkx_labels(self.network, pos=layout,font_size=3)

        plt.show()

        ## Accidently deleted the save option; need to re-add
        if save:
            pass
        
    def extract(self, path_dir, target):
        
        target = target.split('.')[0]

        ## Create copy of existing dataframe
        df_copy = self.object.copy()

        ## Check to make sure unique operation exists in the target
        IS_IN_TARGET = df_copy[df_copy.targetIDX == 1].Oper.unique()
        x = [ch.isupper() for ch in IS_IN_TARGET]
        new_labels = IS_IN_TARGET[np.logical_not(x)]
        
        ######## include check if new_labels is empty #########

        ## Pop off the labels 
        labels = df_copy.pop('targetIDX')
        
        ## Check if //Datasets exists, if not create and modify path
        if not os.path.exists('Datasets'):
            try:
                os.mkdir(path_dir + '//Datasets')
                path_dir_data = os.path.join(path_dir, 'Datasets')
            except OSError:
                print("Failed to create directory")
            else:
                print(f"Created directory at {path_dir_data}") 

        ## Loop through the operations and keep track of index
        for num, oper in enumerate(new_labels):

            ## Create df object of the current operation
            df_oper = df_copy[df_copy["Oper"] == oper]

            dct = defaultdict(list)

            ## Cehck to make sure a Detail header exists
            try:
                ## Loop through the Detail header (idx, details(head, value))
                for idx, item in df_oper['Detail'].items():
                    try:
                        [dct[head].append(value) for head, value in [x.split(': ', maxsplit=1) for x in re.split(r', (?=[^\s:]+\:)', item)]]   
                        dct['index'].append(idx)
                        dct['targetIDX'].append(labels[idx])
                        
                    except ValueError:
                        continue

                    new_df = pd.DataFrame(
                                dict([(k, pd.Series(v, dtype=object)) for k, v in dct.items()])).set_index('index', inplace=False)

                    new_df = pd.concat([new_df, df_oper.drop('Detail')], axis=1)
                        
            except:
                new_df = df_oper
                
            if not os.path.exists(path_dir_data + '//' + target):
                try:
                    os.mkdir(path_dir_data + '//' + target)
                    path_dir_target = os.path.join(path_dir_data, target)
                except OSError:
                    print("Failed to create directory")
                else:
                    print(f"Created directory at {path_dir}")  
                
            if os.path.exists(os.path.join(path_dir_target, oper + '.csv')):
                print("File already exists; appending")   
                new_df.to_csv(os.path.join(path_dir_target, oper + '.csv'), mode='a', encoding='utf-8', index=True, compression=None, header=False)
            else:
                new_df.to_csv(os.path.join(path_dir_target, oper + '.csv'), encoding='utf-8', index=True, compression=None)
                print(f'Saving {oper} to //Datasets')


            
    def _columns(self):
        print("Deprecated since version 0.4. Please use object._info")
        return self.object.columns

    def _info(self):
        n_samples, n_features = self.object.shape
        print(f"Dataset contains {n_samples} unique entries, with {n_features} \
              features. Including {[i for i in self.object.columns]}")
        return None

    def _return_df(self):
        return self.object

    def _unique(self, group_by):
        return self.object[group_by].unique()


def create_color_map(G, attrib, sb_palette="RdBu_r"):

    attributes = [G.node[label][attrib] for label in G.nodes()]

    attributes_unique = list(set(attributes))
    num_values = len(attributes_unique)

    palette = sb.color_palette(sb_palette, num_values).as_hex()
    color_map = dict(zip(attributes_unique, palette))
    node_colors = [color_map[attribute] for attribute in attributes]

    return node_colors


import networkx as nx
import matplotlib.pyplot as plt
import matplotlib as mpl
import pandas as pd
import seaborn as sb
import numpy as np
import re
import os

from collections import defaultdict
from Sandy.pp import importProc

class Network:

    def __init__(self, PATH, import_type='mod', compress=False, track_process=False, verbose=False):
        
        assert import_type in {'mod', 'event'}, f"Import type not supported, choose from 'mod' or 'event'; you chose {import_type}"
        
        if import_type == 'mod':
            self.object = importProc.import_mod(PATH, transform=True, compress=compress, track_process=track_process)
        else:
            self.object = importProc.import_event(PATH, transform=True, compress=compress, track_process=track_process, verbose=verbose)
            
        self.PATH = PATH
        self.using_events = False if import_type == 'mod' else True 

    def produce_edge_list(self, source, target, target_node=False):

        self.network = nx.from_pandas_edgelist(self.object,
                                               source=source,
                                               target=target,
                                               edge_attr=True)

    def assign_labels(self, label_name='ProcName'):

        labels = dict(zip(self.network.nodes(), self.object[label_name]))
        self.network = nx.relabel_nodes(self.network, labels, copy=False)

    def assign_attribute(self, attr_name='targetIDX'):

        attr = dict(zip(self.network.nodes(), self.object[attr_name]))
        nx.set_node_attributes(self.network, attr, 'NodeAttr')

        self.attr_colors = create_color_map(self.network, 'NodeAttr')

    def create_target(self, target_name, new_node=False, verbose=False):
        
        if self.using_events:
            from Sandy.pp import importProc
            object = importProc.import_mod(self.PATH, transform=True, compress=False, track_process=False)
        else:
            object = self.object

        # Extract the PID of the parent process
        target = object.loc[object['ProcName'] == target_name, ['PID']]

        target_len = len(target.PID.unique())

        while True:

            target = object.loc[object.PPID.isin(target.PID.unique()) | \
                            object.PID.isin(target.PID.unique()), ['PID']]
            
            # If length stays the same size, the malicious tree must be covered
            if target_len == len(target.PID.unique()):
                
                object = object.assign(targetIDX = object.PID.isin(target.PID).astype(int))
                
                self.target_PID = target.PID.unique()
                self.target_process = object.iloc[target.index].ProcName.unique()

                if self.using_events:
                    self.object = self.object.assign(targetIDX = (self.object).PID.isin(target.PID).astype(int))
                else:
                    self.object = object
                    
                break
          
            else: ## Set new len to old len
                target_len = len(target.PID.unique())
                
        if verbose:
            print(f"Found {((self.object).targetIDX == 1).sum()} target entries")
            print(f"Unique PID for chosen target at: {self.target_PID}")
            print(f"Unique Process Names for chosen target at: {self.target_process}")
            #print(f"Unique Process Names for chosen target at: #['8fc63e8c5de9e771badefaa50cc2f7a0', 'file-2640385_exe']")

        ## Need to work on graphing. As of right now create target does not produce self.network object properly
        if new_node:
            return Network(self.object, self.PATH)
        
        ## Need to make sure it returns node when I need it to
        ## if new_node:
        #net = Network(self.object, self.PATH)
        #net.produce_edge_list("PID", "PPID")
        
    def plot_network(self, plot_attr='targetIDX', save=False, map_feature=False):

        from itertools import count

        fig = plt.figure()
        layout = nx.spring_layout(self.network)

        ## Default to targetIDX if an attribute doesn't exist?
        ## if targetIDX doesnt exist, add it
        #self.assign_attribute('targetIDX')

        ## Retreive create times and set to line widths
        #durations = [int(i[plot_attr]) for i in dict(self.network.edges).values()]
        #min_, max_ = min(durations), max(durations)
        #durations = [(0.1 - 1)*((i - max_)/(min_ - max_)) + 1  for i in durations]

        ## Set node colors to the targetIDX
        groups = set(nx.get_node_attributes(self.network, plot_attr).values())
        mapping = dict(zip(groups, count()))
        nodes = self.network.nodes()
        colors = [mapping[self.network.nodes[n][plot_attr]] for n in nodes]

        nx.draw_networkx_nodes(self.network,pos=layout,node_size=10,
                               node_color=colors,node_list=nodes)
        nx.draw_networkx_edges(self.network,pos=layout,width=durations,style='dashed',
                               edge_color='red',edge_cmap=plt.cm.Blues)
        nx.draw_networkx_labels(self.network, pos=layout,font_size=3)

        plt.show()

        ## Accidently deleted the save option; need to re-add
        if save:
            pass
        
    def extract(self, path_dir, target):
        
        target = target.split('.')[0]

        ## Create copy of existing dataframe
        df_copy = self.object.copy()

        ## Check to make sure unique operation exists in the target
        IS_IN_TARGET = df_copy[df_copy.targetIDX == 1].Oper.unique()
        x = [ch.isupper() for ch in IS_IN_TARGET]
        new_labels = IS_IN_TARGET[np.logical_not(x)]
        
        ######## include check if new_labels is empty #########

        ## Pop off the labels 
        labels = df_copy.pop('targetIDX')
        
        ## Check if //Datasets exists, if not create and modify path
        if not os.path.exists('Datasets'):
            try:
                os.mkdir(path_dir + '//Datasets')
                path_dir_data = os.path.join(path_dir, 'Datasets')
            except OSError:
                print("Failed to create directory")
            else:
                print(f"Created directory at {path_dir_data}") 

        ## Loop through the operations and keep track of index
        for num, oper in enumerate(new_labels):

            ## Create df object of the current operation
            df_oper = df_copy[df_copy["Oper"] == oper]

            dct = defaultdict(list)

            ## Cehck to make sure a Detail header exists
            try:
                ## Loop through the Detail header (idx, details(head, value))
                for idx, item in df_oper['Detail'].items():
                    try:
                        [dct[head].append(value) for head, value in [x.split(': ', maxsplit=1) for x in re.split(r', (?=[^\s:]+\:)', item)]]   
                        dct['index'].append(idx)
                        dct['targetIDX'].append(labels[idx])
                        
                    except ValueError:
                        continue

                    new_df = pd.DataFrame(
                                dict([(k, pd.Series(v, dtype=object)) for k, v in dct.items()])).set_index('index', inplace=False)

                    new_df = pd.concat([new_df, df_oper.drop('Detail')], axis=1)
                        
            except:
                new_df = df_oper
                
            if not os.path.exists(path_dir_data + '//' + target):
                try:
                    os.mkdir(path_dir_data + '//' + target)
                    path_dir_target = os.path.join(path_dir_data, target)
                except OSError:
                    print("Failed to create directory")
                else:
                    print(f"Created directory at {path_dir}")  
                
            if os.path.exists(os.path.join(path_dir_target, oper + '.csv')):
                print("File already exists; appending")   
                new_df.to_csv(os.path.join(path_dir_target, oper + '.csv'), mode='a', encoding='utf-8', index=True, compression=None, header=False)
            else:
                new_df.to_csv(os.path.join(path_dir_target, oper + '.csv'), encoding='utf-8', index=True, compression=None)
                print(f'Saving {oper} to //Datasets')


            
    def _columns(self):
        print("Deprecated since version 0.4. Please use object._info")
        return self.object.columns

    def _info(self):
        n_samples, n_features = self.object.shape
        print(f"Dataset contains {n_samples} unique entries, with {n_features} \
              features. Including {[i for i in self.object.columns]}")
        return None

    def _return_df(self):
        return self.object

    def _unique(self, group_by):
        return self.object[group_by].unique()


def create_color_map(G, attrib, sb_palette="RdBu_r"):

    attributes = [G.node[label][attrib] for label in G.nodes()]

    attributes_unique = list(set(attributes))
    num_values = len(attributes_unique)

    palette = sb.color_palette(sb_palette, num_values).as_hex()
    color_map = dict(zip(attributes_unique, palette))
    node_colors = [color_map[attribute] for attribute in attributes]

    return node_colors

# Extract the PID of the parent process # If length stays the same size, the malicious tree must be covered ## Set new len to old len #print(f"Unique Process Names for chosen target at: #['8fc63e8c5de9e771badefaa50cc2f7a0', 'file-2640385_exe']") ## Need to work on graphing. As of right now create target does not produce self.network object properly ## Need to make sure it returns node when I need it to ## if new_node: #net = Network(self.object, self.PATH) #net.produce_edge_list("PID", "PPID") ## Default to targetIDX if an attribute doesn't exist? ## if targetIDX doesnt exist, add it #self.assign_attribute('targetIDX') ## Retreive create times and set to line widths #durations = [int(i[plot_attr]) for i in dict(self.network.edges).values()] #min_, max_ = min(durations), max(durations) #durations = [(0.1 - 1)*((i - max_)/(min_ - max_)) + 1  for i in durations] ## Set node colors to the targetIDX ## Accidently deleted the save option; need to re-add ## Create copy of existing dataframe ## Check to make sure unique operation exists in the target ######## include check if new_labels is empty ######### ## Pop off the labels ## Check if //Datasets exists, if not create and modify path ## Loop through the operations and keep track of index ## Create df object of the current operation ## Cehck to make sure a Detail header exists ## Loop through the Detail header (idx, details(head, value))