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

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import os

directory = 'screenshots/seperate/'

initial_val = ''
time_set = []
machine_set = []

for i in sorted(os.listdir(directory)):
    full_path = directory + i
    split_full_path = full_path.split('_')
    print (split_full_path)
    initial_val = split_full_path[0]
    time_set.append(split_full_path[1])
    machine_set.append('_'.join(split_full_path[2:]))
    #img = plt.imread(directory + i)
    #plt.imshow(img)
    #plt.show()

time_set = sorted(list(set(time_set)))
machine_set = sorted(list(set(machine_set)))
#print (initial_val)
#print (time_set)
#print (machine_set)

#for i in time_set:
#    d = initial_val + '_' + i
#    print (d)
#
#    img1 = plt.imread(d + '_' + machine_set[0])
#    img2 = plt.imread(d + '_' + machine_set[1])
#    img3 = plt.imread(d + '_' + machine_set[2])
#    img4 = plt.imread(d + '_' + machine_set[3])
#
#    output = np.vstack( [np.hstack([img1, img2]), np.hstack([img3, img4])] )
#
#    plt.imshow(output)
#    #plt.show()
#    fp = 'screenshots/combined/' + i + '.png'
#    print ("save to: " + fp)
#    plt.imsave(fp, output)

pd.set_option('display.max_columns', None)

data = pd.read_csv('vboxmanage.csv', header=None)
#data.columns = data.iloc[0].values
#data.columns = 'dt'

m = [2,20,38,56]

plt.figure(figsize=(20,10))
node1, = plt.plot(data[m[0]], label='Node 1')
node2, = plt.plot(data[m[1]], label='Node 2')
node3, = plt.plot(data[m[2]], label='Node 3')
node4, = plt.plot(data[m[3]], label='Node 4')
plt.legend(handles=[node1, node2, node3, node4])
plt.title("CPU_User")
plt.show()

pnt = 6
plt.figure(figsize=(20,10))
plt.plot(data[m[0]+pnt])
plt.plot(data[m[1]+pnt])
plt.plot(data[m[2]+pnt])
plt.plot(data[m[3]+pnt])
plt.title("RAM Used")
plt.show()

pnt = 9
plt.figure(figsize=(20,10))
plt.plot(data[m[0]+pnt])
plt.plot(data[m[1]+pnt])
plt.plot(data[m[2]+pnt])
plt.plot(data[m[3]+pnt])
plt.title("Disk Used")
plt.show()



# import pyshark
# cap = pyshark.FileCapture('./12-08-2020/12-08-2020-11-58-14.pcapng')
# dir(cap)

# Export from Wireshark as CSV (easier than dealing with pcap in Python - although PyShark does work for this)
data = pd.read_csv('wireshark.csv')

# SMB START
SMB_data = data[data['Protocol'] =='SMB']
unique_IPs = list(SMB_data['Source'].unique())
unique_IPs2 = list(SMB_data['Destination'].unique())
pnts_src = []
pnts_dst = []

color_set = ['orange', 'blue', 'purple']

for i in range(SMB_data.shape[0]):
    ddd = SMB_data.iloc[i]
    pnts_src.append([i, unique_IPs.index(ddd['Source']), ddd['Length']])
    pnts_dst.append([i, unique_IPs2.index(ddd['Destination']), ddd['Length']])
    #print (i, ddd['Src'], '-->', ddd['Dest'], ddd['Length'])
    #plt.scatter(i, unique_IPs.index(ddd['Src']), c='g')
    #plt.scatter(i, unique_IPs.index(ddd['Dest']), c='r')
#plt.show()

pnts_src = np.array(pnts_src)
pnts_dst = np.array(pnts_dst)
plt.figure(figsize=(100,10))
plt.scatter(pnts_src[:,0], pnts_src[:,1], c='blue', s=pnts_src[:,2])
plt.scatter(pnts_dst[:,0], pnts_dst[:,1], c='orange', s=pnts_dst[:,2])
plt.plot(pnts_src[:,0], pnts_src[:,1], pnts_dst[:,0], pnts_dst[:,1], color='purple', alpha=0.5)
blue_patch = mpatches.Patch(color='blue', label='Source')
orange_patch = mpatches.Patch(color='orange', label='Destination')
purple_patch = mpatches.Patch(color='purple', label='SMB')
plt.yticks(np.arange(4), ('Node 1', 'Node 2', 'Node 3', 'Node 4'))
plt.xticks([])
plt.legend(handles=[blue_patch, orange_patch, purple_patch])
plt.title("SMB Activity")
plt.show()
# SMB END

# TCP START
SMB_data = data[data['Protocol'] =='TCP']
unique_IPs = list(SMB_data['Source'].unique())
print(unique_IPs)
pnts_src = []
pnts_dst = []

color_set = ['orange', 'blue', 'purple']

for i in range(SMB_data.shape[0]):
    ddd = SMB_data.iloc[i]
    pnts_src.append([i, unique_IPs.index(ddd['Source']), ddd['Length']])
    pnts_dst.append([i, unique_IPs.index(ddd['Destination']), ddd['Length']])
    #print (i, ddd['Src'], '-->', ddd['Dest'], ddd['Length'])
    #plt.scatter(i, unique_IPs.index(ddd['Src']), c='g')
    #plt.scatter(i, unique_IPs.index(ddd['Dest']), c='r')
#plt.show()

pnts_src = np.array(pnts_src)
pnts_dst = np.array(pnts_dst)
plt.figure(figsize=(100,10))
plt.scatter(pnts_src[:,0], pnts_src[:,1], c='b', s=pnts_src[:,2])
plt.scatter(pnts_dst[:,0], pnts_dst[:,1], c='orange', s=pnts_dst[:,2])
#plt.plot(pnts_src[:,0], pnts_src[:,1], 'b--', pnts_dst[:,0], pnts_dst[:,1], alpha=0.5)
plt.plot(pnts_src[:,0], pnts_src[:,1], pnts_dst[:,0], pnts_dst[:,1], color='purple', alpha=0.5)
blue_patch = mpatches.Patch(color='blue', label='Source')
orange_patch = mpatches.Patch(color='orange', label='Destination')
purple_patch = mpatches.Patch(color='purple', label='TCP')
plt.yticks(np.arange(len(unique_IPs)), unique_IPs)
plt.xticks([])
plt.legend(handles=[blue_patch, orange_patch, purple_patch])
plt.title("TCP Activity")
plt.show()
# TCP END

# SSDP START
SMB_data = data[data['Protocol'] =='SSDP']
unique_IPs = list(SMB_data['Source'].unique()) + list(SMB_data['Destination'].unique())
unique_IPs = list(dict.fromkeys(unique_IPs))

pnts_src = []
pnts_dst = []

color_set = ['orange', 'blue', 'purple']

for i in range(SMB_data.shape[0]):
    ddd = SMB_data.iloc[i]
    pnts_src.append([i, unique_IPs.index(ddd['Source']), ddd['Length']])
    pnts_dst.append([i, unique_IPs.index(ddd['Destination']), ddd['Length']])
    #print (i, ddd['Src'], '-->', ddd['Dest'], ddd['Length'])
    #plt.scatter(i, unique_IPs.index(ddd['Src']), c='g')
    #plt.scatter(i, unique_IPs.index(ddd['Dest']), c='r')
#plt.show()

pnts_src = np.array(pnts_src)
pnts_dst = np.array(pnts_dst)
plt.figure(figsize=(100,10))
plt.scatter(pnts_src[:,0], pnts_src[:,1], c='b', s=pnts_src[:,2])
plt.scatter(pnts_dst[:,0], pnts_dst[:,1], c='orange', s=pnts_dst[:,2])
#plt.plot(pnts_src[:,0], pnts_src[:,1], 'b--', pnts_dst[:,0], pnts_dst[:,1], alpha=0.5)
plt.plot(pnts_src[:,0], pnts_src[:,1], pnts_dst[:,0], pnts_dst[:,1], color='purple', alpha=0.5)
blue_patch = mpatches.Patch(color='blue', label='Source')
orange_patch = mpatches.Patch(color='orange', label='Destination')
purple_patch = mpatches.Patch(color='purple', label='TCP')
plt.yticks(np.arange(len(unique_IPs)), unique_IPs)
plt.xticks([])
plt.legend(handles=[blue_patch, orange_patch, purple_patch])
plt.title("SSDP Activity")
plt.show()
# SSDP END

# UDP START
SMB_data = data[data['Protocol'] =='UDP']
unique_IPs = list(SMB_data['Source'].unique()) + list(SMB_data['Destination'].unique())
yaxis = list(dict.fromkeys(unique_IPs))

pnts_src = []
pnts_dst = []

color_set = ['orange', 'blue', 'purple']

for i in range(SMB_data.shape[0]):
    ddd = SMB_data.iloc[i]
    pnts_src.append([i, unique_IPs.index(ddd['Source']), ddd['Length']])
    pnts_dst.append([i, unique_IPs.index(ddd['Destination']), ddd['Length']])
    #print (i, ddd['Src'], '-->', ddd['Dest'], ddd['Length'])
    #plt.scatter(i, unique_IPs.index(ddd['Src']), c='g')
    #plt.scatter(i, unique_IPs.index(ddd['Dest']), c='r')
#plt.show()

pnts_src = np.array(pnts_src)
pnts_dst = np.array(pnts_dst)
plt.figure(figsize=(100,10))
plt.scatter(pnts_src[:,0], pnts_src[:,1], c='b', s=pnts_src[:,2])
plt.scatter(pnts_dst[:,0], pnts_dst[:,1], c='orange', s=pnts_dst[:,2])
#plt.plot(pnts_src[:,0], pnts_src[:,1], 'b--', pnts_dst[:,0], pnts_dst[:,1], alpha=0.5)
plt.plot(pnts_src[:,0], pnts_src[:,1], pnts_dst[:,0], pnts_dst[:,1], color='purple', alpha=0.5)
blue_patch = mpatches.Patch(color='blue', label='Source')
orange_patch = mpatches.Patch(color='orange', label='Destination')
purple_patch = mpatches.Patch(color='purple', label='SSDP')
plt.yticks(np.arange(len(yaxis)), yaxis)
plt.xticks([])
plt.legend(handles=[blue_patch, orange_patch, purple_patch])
plt.title("UDP Activity")
plt.show()
# UDP END

# BROWSER START
SMB_data = data[data['Protocol'] =='BROWSER']
unique_IPs = list(SMB_data['Source'].unique()) + list(SMB_data['Destination'].unique())
#yaxis = list(SMB_data['Source'].unique())

pnts_src = []
pnts_dst = []

color_set = ['orange', 'blue', 'purple']

for i in range(SMB_data.shape[0]):
    ddd = SMB_data.iloc[i]
    pnts_src.append([i, unique_IPs.index(ddd['Source']), ddd['Length']])
    pnts_dst.append([i, unique_IPs.index(ddd['Destination']), ddd['Length']])
    #print (i, ddd['Src'], '-->', ddd['Dest'], ddd['Length'])
    #plt.scatter(i, unique_IPs.index(ddd['Src']), c='g')
    #plt.scatter(i, unique_IPs.index(ddd['Dest']), c='r')
#plt.show()

pnts_src = np.array(pnts_src)
pnts_dst = np.array(pnts_dst)
plt.figure(figsize=(100,10))
plt.scatter(pnts_src[:,0], pnts_src[:,1], c='b', s=pnts_src[:,2])
plt.scatter(pnts_dst[:,0], pnts_dst[:,1], c='orange', s=pnts_dst[:,2])
#plt.plot(pnts_src[:,0], pnts_src[:,1], 'b--', pnts_dst[:,0], pnts_dst[:,1], alpha=0.5)
plt.plot(pnts_src[:,0], pnts_src[:,1], pnts_dst[:,0], pnts_dst[:,1], color='purple', alpha=0.5)
blue_patch = mpatches.Patch(color='blue', label='Source')
orange_patch = mpatches.Patch(color='orange', label='Destination')
purple_patch = mpatches.Patch(color='purple', label='BROWSER')
plt.yticks(np.arange(len(unique_IPs)), unique_IPs)
plt.xticks([])
plt.legend(handles=[blue_patch, orange_patch, purple_patch])
plt.title("BROWSER Activity")
plt.show()

data['Protocol'].unique()
# BROWSER END

# HTTP START
SMB_data = data[data['Protocol'] =='HTTP']
unique_IPs = list(SMB_data['Source'].unique()) + list(SMB_data['Destination'].unique())
yaxis = list(SMB_data['Source'].unique())

pnts_src = []
pnts_dst = []

color_set = ['orange', 'blue', 'purple']

for i in range(SMB_data.shape[0]):
    ddd = SMB_data.iloc[i]
    pnts_src.append([i, unique_IPs.index(ddd['Source']), ddd['Length']])
    pnts_dst.append([i, unique_IPs.index(ddd['Destination']), ddd['Length']])
    #print (i, ddd['Src'], '-->', ddd['Dest'], ddd['Length'])
    #plt.scatter(i, unique_IPs.index(ddd['Src']), c='g')
    #plt.scatter(i, unique_IPs.index(ddd['Dest']), c='r')
#plt.show()

pnts_src = np.array(pnts_src)
pnts_dst = np.array(pnts_dst)
plt.figure(figsize=(100,10))
plt.scatter(pnts_src[:,0], pnts_src[:,1], c='b', s=pnts_src[:,2])
plt.scatter(pnts_dst[:,0], pnts_dst[:,1], c='orange', s=pnts_dst[:,2])
#plt.plot(pnts_src[:,0], pnts_src[:,1], 'b--', pnts_dst[:,0], pnts_dst[:,1], alpha=0.5)
plt.plot(pnts_src[:,0], pnts_src[:,1], pnts_dst[:,0], pnts_dst[:,1], color='purple', alpha=0.5)
blue_patch = mpatches.Patch(color='blue', label='Source')
orange_patch = mpatches.Patch(color='orange', label='Destination')
purple_patch = mpatches.Patch(color='purple', label='HTTP')
plt.yticks(np.arange(len(yaxis)), yaxis)
plt.xticks([])
plt.legend(handles=[blue_patch, orange_patch, purple_patch])
plt.title("HTTP Activity")
plt.show()
# HTTP END

# SMB PIPE START
SMB_data = data[data['Protocol'] =='SMB Pipe']
unique_IPs = list(SMB_data['Source'].unique()) + list(SMB_data['Destination'].unique())
yaxis = list(SMB_data['Source'].unique())
pnts_src = []
pnts_dst = []

color_set = ['orange', 'blue', 'purple']

for i in range(SMB_data.shape[0]):
    ddd = SMB_data.iloc[i]
    pnts_src.append([i, unique_IPs.index(ddd['Source']), ddd['Length']])
    pnts_dst.append([i, unique_IPs.index(ddd['Destination']), ddd['Length']])
    #print (i, ddd['Src'], '-->', ddd['Dest'], ddd['Length'])
    #plt.scatter(i, unique_IPs.index(ddd['Src']), c='g')
    #plt.scatter(i, unique_IPs.index(ddd['Dest']), c='r')
#plt.show()

pnts_src = np.array(pnts_src)
pnts_dst = np.array(pnts_dst)
plt.figure(figsize=(100,10))
plt.scatter(pnts_src[:,0], pnts_src[:,1], c='b', s=pnts_src[:,2])
plt.scatter(pnts_dst[:,0], pnts_dst[:,1], c='orange', s=pnts_dst[:,2])
#plt.plot(pnts_src[:,0], pnts_src[:,1], 'b--', pnts_dst[:,0], pnts_dst[:,1], alpha=0.5)
plt.plot(pnts_src[:,0], pnts_src[:,1], pnts_dst[:,0], pnts_dst[:,1], color='purple', alpha=0.5)
blue_patch = mpatches.Patch(color='blue', label='Source')
orange_patch = mpatches.Patch(color='orange', label='Destination')
purple_patch = mpatches.Patch(color='purple', label='SMB Pipe')
plt.yticks(np.arange(len(yaxis)), yaxis)
plt.xticks([])
plt.legend(handles=[blue_patch, orange_patch, purple_patch])
plt.title("SMB Pipe Activity")
plt.show()
# SMB PIPE END
