import matplotlib.pyplot as plt
import numpy as np
import json
import os
from IPython.display import Markdown, display, HTML
from matplotlib.pyplot import *
import pandas as pd
import seaborn as sb
from matplotlib.lines import Line2D

# proměnná na scaling osy x
pole = []
for i in range(0, 256):
    pole.append(i)

def set_graph_params(*args):

    # parsing parametrů
    if len(args)==2:
        return {
            "label1" : args[0],
            "label2": args[1]
    }
    if len(args)==3:
        return {
            "x" : args[0],
            "title": args[1],
            "colors": args[2]
    }
    if len(args)==4:
        return {
            "x" : args[0],
            "y" : args[1],
            "title" : args[2],
            "colors": args[3]
    }
    if len(args)==5:
        return {
            "x" : args[0],
            "y" : args[1],
            "title" : args[2],
            "label": args[3],
            "colors": args[4],
        }
    if len(args)==6:
        return {
            "x" : args[0],
            "y" : args[1],
            "title" : args[2],
            "label1": args[3],
            "label2": args[4],
            "colors": args[5],
        }

def protocol_determination(data):

    #zjištění užitého protokolu na základě cílového portu
    if (data['pr'] == 6) & (data['dp'] == 443):
        protocol = 'HTTPS'
    elif (data['pr'] == 6) & (data['dp'] == 80):
        protocol = 'HTTP'
    elif (data['pr'] == 6) & (data['dp'] == 22):
        protocol = 'SSH'
    elif (data['pr'] == 6) & (data['dp'] == 465):
        protocol = 'SMTPS'
    else:
        protocol = 'unknown'
    return protocol

def load_data(file_names,final_dir):

    # proměnné a inicializace
    file_count = len(file_names)
    l=0
    bytes_out = []
    bytes_in = []
    packets_out = []
    packets_in = []
    flow_time = []
    avg_byte_dist = []
    byte_seq = []
    time_seq = []
    dir_seq = []
    splt_p = []
    splt_t = []
    splt_d = []
    srlt_p = []
    srlt_t = []
    srlt_d = []
    labels = []
    for i in range(0, 256):
        avg_byte_dist.append(0)

    # načtení dat
    for key,value in file_names.items():
        with open (final_dir+'\\'+key) as json_file:
            records = json_file.readlines()
        try:
            data = json.loads(records[1])
        except IndexError:
            data = json.loads(records[0])
        # základní data
        bytes_out.append(data['bytes_out'])
        bytes_in.append(data['bytes_in'])
        packets_out.append(data['num_pkts_out'])
        packets_in.append(data['num_pkts_in'])
        flow_time.append(data['time_end'] - data['time_start'])
        loaded_dist = data['byte_dist']
        for k in loaded_dist[0:255]:
            avg_byte_dist[l] += k
            l += 1
        l = 0

        # načtení SPLT dat
        splt=data['packets']
        for record in splt:
            try:
                byte_seq.append(record['b'])
                time_seq.append(record['ipt'])
                dir_seq.append(record['dir'])
            except KeyError:
                pass
        splt_p.append(byte_seq[:])
        splt_t.append(time_seq[:])
        splt_d.append(dir_seq[:])
        byte_seq.clear()
        time_seq.clear()
        dir_seq.clear()

        # načtení SRLT dat, pokud jsou přítomna
        if ('tls' in data):
            if ('srlt' in data['tls']):
                labels.append(value)
                srlt=data['tls']['srlt']
                for record in srlt:
                    try:
                        byte_seq.append(record['b'])
                        time_seq.append(record['ipt'])
                        dir_seq.append(record['dir'])
                    except KeyError:
                        pass
                srlt_p.append(byte_seq[:])
                srlt_t.append(time_seq[:])
                srlt_d.append(dir_seq[:])
                byte_seq.clear()
                time_seq.clear()
                dir_seq.clear()

    # zprůměrování distribuce bajtů
    for i in range(0, len(avg_byte_dist)):
        avg_byte_dist[i] = avg_byte_dist[i] / file_count

    # uložení načtených dat do dict a jeho vrácení
    loaded = {
        "bytes_in" : bytes_in,
        "bytes_out" : bytes_out,
        "packets_in" : packets_in,
        "packets_out" : packets_out,
        "flow_time" : flow_time,
        "avg_byte_dist" : avg_byte_dist,
        "splt_t" : splt_t,
        "splt_p" : splt_p,
        "splt_d" : splt_d,
        "labels" : labels,
        "srlt_t": srlt_t,
        "srlt_p": srlt_p,
        "srlt_d": srlt_d,
    }
    return loaded

def bytes_in_out(loaded,params,file_names):

    # získání parametrů
    notes = file_names.values()
    file_count = len(file_names)
    x = params.get("x")
    y = params.get("y")
    header = params.get("title")
    colors = params.get("colors")
    bytes_in=loaded.get("bytes_in")
    bytes_out=loaded.get("bytes_out")

    # parametry grafu
    index = np.arange(file_count)
    bar_width = 0.35
    opacity = 0.8
    plt.rcParams['figure.dpi'] = 150
    plt.subplots()

    # vykreslení grafů
    plt.bar(index, bytes_in, bar_width,
                     alpha=opacity,
                     color=colors[0],
                     label='in')

    plt.bar(index + bar_width, bytes_out, bar_width,
                     alpha=opacity,
                     color=colors[1],
                     label='out')

    # popisky a legenda
    plt.xlabel(x)
    plt.ylabel(y)
    plt.title(header)
    plt.legend()
    if (file_count<=4):
        plt.xticks(index + bar_width/2, notes)
    else:
        plt.xticks(index + bar_width/2,notes, rotation=90)
    plt.show()

def packets_in_out(loaded,params,file_names):

    # získání parametrů
    notes = file_names.values()
    file_count = len(file_names)
    x = params.get("x")
    y = params.get("y")
    header = params.get("title")
    colors = params.get("colors")
    packets_in = loaded.get("packets_in")
    packets_out = loaded.get("packets_out")

    # parametry grafu
    index = np.arange(file_count)
    bar_width = 0.35
    opacity = 0.8
    plt.rcParams['figure.dpi'] = 150
    plt.subplots()

    # vykreslení grafů
    plt.bar(index, packets_in, bar_width,
                     alpha=opacity,
                     color=colors[0],
                     label='in')
    plt.bar(index + bar_width, packets_out, bar_width,
                     alpha=opacity,
                     color=colors[1],
                     label='out')

    # popisky a legenda
    plt.xlabel(x)
    plt.ylabel(y)
    plt.title(header)
    plt.legend()
    if (file_count <= 4):
        plt.xticks(index + bar_width / 2, notes)
    else:
        plt.xticks(index + bar_width / 2, notes, rotation=90)
    plt.show()

def flow_duration(loaded,params,file_names):
    # získání parametrů
    notes = file_names.values()
    file_count = len(file_names)
    x = params.get("x")
    y = params.get("y")
    header = params.get("title")
    label = params.get("label")
    colors = params.get("colors")
    flow_time=loaded.get("flow_time")

    #pokud je vybrán jeden soubor, zobrazí se pouze výpis namísto grafu
    if (file_count==1):
        display(Markdown("<h3>IP tok trval "+str(flow_time[0])+" vteřin.</h3>"))
    else:
        # parametry grafu
        index = np.arange(file_count)
        bar_width = 0.35
        opacity = 0.8
        plt.rcParams['figure.dpi'] = 150
        plt.subplots()

        # vykreslení grafu
        plt.bar(index, flow_time, bar_width,
                         alpha=opacity,
                         color=colors[0],
                         label=label)

        # popisky a legenda
        plt.xlabel(x)
        plt.ylabel(y)
        plt.title(header)
        plt.legend()
        if (file_count <= 4):
            plt.xticks(index, notes)
        else:
            plt.xticks(index, notes, rotation=90)
        plt.show()

def byte_dist_load_one(byte_dist_file, final_dir):

    #načtení dat jednoho souboru specificky pro porovnání a vrácení pole s hodnotami
    with open (final_dir+'\\'+byte_dist_file) as json_file:
        records = json_file.readlines()
    try:
        data = json.loads(records[1])
    except IndexError:
        data = json.loads(records[0])
    byte_dist=data['byte_dist']
    return byte_dist

def byte_dist_show(byte_dist_file,params,final_dir):

    # zobrazení grafu distribuce bajtů
    x = params.get("x")
    y = params.get("y")
    header = params.get("title")
    colors = params.get("colors")
    plt.bar(pole, byte_dist_load_one(byte_dist_file,final_dir),color=colors[0])
    plt.xlabel(x)
    plt.ylabel(y)
    plt.title(header)
    plt.show()

def byte_dist_compare(loaded,params,byte_dist_file,final_dir):

    # získání parametrů
    x = params.get("x")
    y = params.get("y")
    header = params.get("title")
    file_label = params.get("label1")
    avg_label = params.get("label2")
    colors = params.get("colors")
    avg_byte_dist=loaded.get("avg_byte_dist")

    # načtení jednoho souboru k porovnání
    chosen_byte_dist=byte_dist_load_one(byte_dist_file, final_dir)

    #normalizace hodnot
    chosen_byte_dist = normalize(chosen_byte_dist)
    avg_byte_dist = normalize(avg_byte_dist)

    plt.subplots()

    # vykreslení grafů
    plt.bar(pole,chosen_byte_dist,color=colors[0],alpha=0.5,label=file_label)
    plt.bar(pole,avg_byte_dist,color=colors[1],alpha=0.6,label=avg_label)

    # popisky a legenda
    plt.xlabel(x)
    plt.ylabel(y)
    plt.title(header)
    plt.legend()
    plt.show()

def normalize(array_to_normalize):

    # převod na procenta
    total = 0
    result = []
    for i in range(0, len(array_to_normalize)):
        total = total + array_to_normalize[i]
        result.append(0)
    for i in range(0, len(array_to_normalize)):
        result[i] = (array_to_normalize[i]/total)*100
    return result

def used_ciphersuite(cipher_file,final_dir):

    # načtení dat
    cs = ''
    with open (final_dir+'\\'+cipher_file) as json_file:
        records = json_file.readlines()
    try:
        data = json.loads(records[1])
    except IndexError:
        data = json.loads(records[0])

    # určení užitého protokolu
    protocol=protocol_determination(data)

    # pokud je zachycený tok HTTPS či SMTPS, porovná se vůči datům CS v souboru a vypíše se typ užité CS a její síla
    if (protocol == "HTTPS" or protocol == "SMTPS"):
        scs = data['tls']['scs']
        cs = data['tls']['cs']
        found = False
        with open('lib/ciphersuites.txt') as f:
            for line in f:
                if scs in line:
                    found = True
                    ciphersuite = line.split(",")
                    if str(ciphersuite[2]) == "ok":
                        stav = "silná"
                    elif str(ciphersuite[2]) == "zastarale":
                        stav = "slabá"
                    else:
                        stav = ""
            else:
                if found == False:
                    stav = "unknown"
            if (stav == "silná") | (stav == "slabá"):
                display(Markdown('<h3>Je užita ' + stav + ' cipher suite ' + str(ciphersuite[1]) + '.</h3>'))
            if stav == "unknown":
                display(Markdown('<h3>Cipher suite \"' + str(scs) + '\" nebyla nalezena v seznamu.</h3>'))
    elif (protocol == "HTTP"):
        display(Markdown('<h3>Jedná se o ' + str(protocol) + ' protokol, spojení není šifrováno.</h3>'))
    else:
        display(Markdown('<h3>Záznam neobsahuje informace o cipher suites.</h3>'))

    # výpis všech možností CS, které jsou uvedeny v datovém doku
    ciphers = []
    if (protocol == "HTTPS" or protocol == "SMTPS"):
        with open('lib/ciphersuites.txt') as f:
            for line in f:
                for cipher in cs:
                    ciphersuite = line.split(",")
                    if cipher == ciphersuite[0]:
                        ciphers.append(ciphersuite[1])

        # zobrazení v tabulce
        df = pd.DataFrame(ciphers)
        df.columns = ['Cipher suites']
        heading_properties = [('font-size', '18px'), ('border-style', 'solid'), ('border-width', '1px'),
                              ('text-align', 'center'), ('background-color', 'lightgrey')]
        cell_properties = [('font-size', '14px'), ('border-style', 'solid'), ('border-width', '1px'),
                           ('text-align', 'left')]
        df_style = [dict(selector="th", props=heading_properties),
                   dict(selector="td", props=cell_properties)]
        display(HTML("""
        <style>
        .table {
            display: flex;
            align-items: center;
            text-align: center;
        }
        </style>
        """))
        display(df.style.hide_index().set_table_styles(df_style))

def byte_dist_avg(loaded,params):

    # zobrazení grafu průměrné distribuce bajtů více souborů
    x = params.get("x")
    y = params.get("y")
    header = params.get("title")
    colors = params.get("colors")
    avg_byte_dist = loaded.get("avg_byte_dist")
    bar(pole, avg_byte_dist,color=colors[0])
    plt.xlabel(x)
    plt.ylabel(y)
    plt.title(header)
    plt.show()

def heatmap_dist(loaded,params,byte_dist_file,final_dir):

    # získání titulků grafů
    title1=params.get("label1")
    title2=params.get("label2")

    # načtení dat pro grafy
    heatmap_array = byte_dist_load_one(byte_dist_file,final_dir)
    heatmap_array_avg = loaded.get("avg_byte_dist")

    #nastavení stejného maxima obou grafů, aby byly barvy porovnatelné
    maximum = max(heatmap_array)
    if maximum < max(heatmap_array_avg):
        maximum = max(heatmap_array_avg)

    # tvarování polí
    data = np.array(heatmap_array)
    shape=(16,16)
    heatmap = data.reshape(shape)
    data_avg = np.array(heatmap_array_avg)
    heatmap_avg = data_avg.reshape(shape)

    # zobrazení obou grafů
    ax = sb.heatmap(heatmap, vmin=0,vmax=maximum,cmap="YlGnBu")
    plt.title(title1)
    ax.set_yticks([])
    plt.show()
    ax = sb.heatmap(heatmap_avg, vmin=0,vmax=maximum,cmap="YlGnBu")
    plt.title(title2)
    ax.set_yticks([])
    plt.show()

def load_splt_data(file_names,final_dir):

    byte_seq = []
    time_seq = []
    dir_seq = []
    splt_p = []
    splt_t = []
    splt_d = []

    # načtení dat SPLT
    for key in file_names:
        with open (final_dir+'\\'+key) as json_file:
            records = json_file.readlines()
        try:
            data = json.loads(records[1])
        except IndexError:
            data = json.loads(records[0])
        splt=data['packets']
        for record in splt:
            try:
                byte_seq.append(record['b'])
                time_seq.append(record['ipt'])
                dir_seq.append(record['dir'])
            except KeyError:
                pass
        splt_p.append(byte_seq[:])
        splt_t.append(time_seq[:])
        splt_d.append(dir_seq[:])
        byte_seq.clear()
        time_seq.clear()
        dir_seq.clear()

    # vrácení dat ohledně mezipaketové mezery, velikosti paketů a směru
    loaded = {
    "splt_t" : splt_t,
    "splt_p" : splt_p,
    "splt_d" : splt_d,
    }
    return loaded

def load_srlt_data(file_names,final_dir):

    byte_seq = []
    time_seq = []
    dir_seq = []
    srlt_p = []
    srlt_t = []
    srlt_d = []
    labels = []

    # načtení dat SRLT
    for key,value in file_names.items():
        with open (final_dir+'\\'+key) as json_file:
            records = json_file.readlines()
        try:
            data = json.loads(records[1])
        except IndexError:
            data = json.loads(records[0])
        if ('tls' in data):
            if ('srlt' in data['tls']):
                labels.append(value)
                srlt = data['tls']['srlt']
                for record in srlt:
                    try:
                        byte_seq.append(record['b'])
                        time_seq.append(record['ipt'])
                        dir_seq.append(record['dir'])
                    except KeyError:
                        pass
                srlt_p.append(byte_seq[:])
                srlt_t.append(time_seq[:])
                srlt_d.append(dir_seq[:])
                byte_seq.clear()
                time_seq.clear()
                dir_seq.clear()
        else:
            return

    # vrácení dat ohledně mezipaketové mezery, velikosti paketů a směru
    loaded = {
    "labels": labels,
    "srlt_t": srlt_t,
    "srlt_p": srlt_p,
    "srlt_d": srlt_d,
    }
    return loaded

def sort_splt_data(splt_t):

    # postupné sčítání mezer do časové linie
    j=1
    for i in range(len(splt_t)):
        while j<(len(splt_t[i])):
            splt_t[i][j]+=splt_t[i][j-1]
            j+=1
        j=1
    return splt_t

def direction_colors(colors,splt_d):

    # směry paketů převedeny na sekvenci barev
    j = 0
    color_seq = splt_d[:]
    for i in range(len(color_seq)):
        while j < (len(color_seq[i])):
            if (color_seq[i][j]=='<'):
                color_seq[i][j]=colors[0]
            else:
                color_seq[i][j] = colors[1]
            j += 1
        j = 0
    return color_seq

def scale_linelength(splt_p):

    # scaling vertikálních čar podle maximální délky čáry v grafu
    j=0
    max=0
    for i in range(len(splt_p)):
        while j < (len(splt_p[i])):
            if (splt_p[i][j]>max):
                max=splt_p[i][j]
            j += 1
        j = 0
    j = 0
    for i in range(len(splt_p)):
        while j < (len(splt_p[i])):
            splt_p[i][j]=splt_p[i][j]/max
            j += 1
        j = 0
    return splt_p

def splt_chart(loaded,params,file_names,final_dir):

    # získání parametrů
    file_count = len(file_names)
    x = params.get("x")
    header = params.get("title")
    barvy = params.get("colors")

    # pokud je vybrán pouze jeden soubor, bude zobrazen jiný graf
    if (file_count==1):
        splt_chart_one(params,file_names,final_dir)
    else:
        plt.figure(None, (10, file_count))
        # získání potřebných dat - scaling velikosti čar, sekvence barev podle směrů, časová linka
        splt_t = sort_splt_data(loaded.get("splt_t"))
        colors = direction_colors(barvy,loaded.get("splt_d"))
        splt_p = scale_linelength(loaded.get("splt_p"))

        # vykreslení čar
        for i,label in enumerate(file_names.values()):
            for j in range(len(splt_p[i])):
                plt.vlines(splt_t[i][j],i-splt_p[i][j]/2,i+splt_p[i][j]/2, label=label,alpha=0.5, colors=colors[i][j])

        # popisky a legenda
        plt.yticks(range(len(file_names.values())), labels=file_names.values())
        plt.ylim(-0.5, len(file_names.values())-0.5)
        ax = plt.gca()
        ax.yaxis.set_minor_locator(MultipleLocator(0.5))
        custom_lines = [Line2D([0], [0], alpha=0.5, color=barvy[0], lw=4),
                        Line2D([0], [0], alpha=0.5, color=barvy[1], lw=4)]
        plt.legend(custom_lines, ['in', 'out'])
        plt.grid(which='minor',linestyle=':')
        plt.title(header)
        plt.xlabel(x)
        plt.show()

def splt_chart_one(params,file_names,final_dir):

    # získání a načtení parametrů
    x = params.get("x")
    y = params.get("y")
    header = params.get("title")
    barvy = params.get("colors")
    loaded = load_splt_data(file_names,final_dir)
    splt_p = loaded.get("splt_p")
    splt_t = sort_splt_data(loaded.get("splt_t"))
    colors = direction_colors(barvy,loaded.get("splt_d"))

    # vykreslení grafu
    plt.vlines(splt_t[0],0,splt_p[0],label='label',alpha=0.5,colors=colors[0])

    # popisky a legenda
    custom_lines = [Line2D([0], [0],alpha=0.5, color=barvy[0], lw=4),
                    Line2D([0], [0],alpha=0.5, color=barvy[1], lw=4)]
    plt.legend(custom_lines,['in','out'])
    plt.title(header)
    plt.xlabel(x)
    plt.ylabel(y)
    plt.show()

def srlt_chart(loaded,params,file_names,final_dir):

    # pokud není položka labels prázná, jsou přítomny záznamy obsahující data o SRLT
    labels = loaded.get("labels")
    if len(labels)!=0:
        file_count = len(labels)
        x = params.get("x")
        header = params.get("title")
        barvy = params.get("colors")
        # pokud je vybrán jeden soubor, bude zobrazen jiný graf
        if (file_count==1):
            srlt_chart_one(params,file_names,final_dir)
        else:
            plt.figure(None, (10, file_count))
            # získání potřebných dat - scaling velikosti čar, sekvence barev podle směrů, časová linka
            srlt_t=sort_splt_data(loaded.get("srlt_t"))
            colors = direction_colors(barvy,loaded.get("srlt_d"))
            srlt_p=scale_linelength(loaded.get("srlt_p"))

            # vykreslení čar
            for i,label in enumerate(labels):
                for j in range(len(srlt_p[i])):
                    plt.vlines(srlt_t[i][j],i-srlt_p[i][j]/2,i+srlt_p[i][j]/2, label=label, alpha=0.5,colors=colors[i][j])

            # popisky a legenda
            custom_lines = [Line2D([0], [0],alpha=0.5,color=barvy[0], lw=4),
                            Line2D([0], [0], alpha=0.5,color=barvy[1], lw=4)]
            plt.legend(custom_lines, ['in', 'out'])
            plt.yticks(range(len(labels)), labels=labels)
            plt.ylim(-0.5, len(labels)-0.5)
            ax=plt.gca()
            ax.yaxis.set_minor_locator(MultipleLocator(0.5))
            plt.grid(which='minor',linestyle=':')
            plt.title(header)
            plt.xlabel(x)
            plt.show()
    else:
        display(Markdown('<h3>Nejsou vybrány žádné toky obsahující SRLT záznamy.</h3>'))

def srlt_chart_one(params,file_names,final_dir):

    # načtení dat vybraného souboru
    loaded = load_srlt_data(file_names,final_dir)

    # pokud nejsou přítomna SRLT data v souboru, je vypsán text
    if not loaded:
        display(Markdown('<h3>Není vybrán tok obsahující SRLT záznam.</h3>'))
        return
    else:
        x = params.get("x")
        y = params.get("y")
        header = params.get("title")
        barvy = params.get("colors")
        srlt_p = loaded.get("srlt_p")
        srlt_t = sort_splt_data(loaded.get("srlt_t"))
        colors = direction_colors(barvy,loaded.get("srlt_d"))

        # vykreslení grafu
        plt.vlines(srlt_t[0],0,srlt_p[0],label='label',alpha=0.5,colors=colors[0])

        # popisky a legenda
        custom_lines = [Line2D([0], [0],alpha=0.5, color=barvy[0], lw=4),
                        Line2D([0], [0],alpha=0.5, color=barvy[1], lw=4)]
        plt.legend(custom_lines, ['in', 'out'])
        plt.title(header)
        plt.xlabel(x)
        plt.ylabel(y)
        plt.show()

