#!/usr/bin/env python
#require python 3.6 or above
from preprocessors import *
import tensorflow.contrib.learn as skflow
import pandas as pd
import os
import io
import random
import requests
import numpy as np
from sklearn import metrics
from scipy.stats import zscore
from sklearn.model_selection import train_test_split
from sklearn import metrics
from keras.models import Sequential
from keras.layers.core import Dense, Activation
from keras.callbacks import EarlyStopping
from keras.utils.data_utils import get_file
from keras.models import load_model

try:
    path = get_file('kddcup.data_10_percent.gz', origin='http://kdd.ics.uci.edu/databases/kddcup99/kddcup.data_10_percent.gz')
except:
    print('Error downloading')
    raise

df = pd.read_csv(path, header=None)
rawcs = df
df.dropna(inplace=True,axis=1)
df.columns = ['duration','protocol_type','service','flag','src_bytes','dst_bytes','land','wrong_fragment','urgent','hot','num_failed_logins','logged_in','num_compromised','root_shell','su_attempted','num_root','num_file_creations','num_shells','num_access_files','num_outbound_cmds','is_host_login','is_guest_login','count','srv_count','serror_rate','srv_serror_rate','rerror_rate','srv_rerror_rate','same_srv_rate','diff_srv_rate','srv_diff_host_rate','dst_host_count','dst_host_srv_count','dst_host_same_srv_rate','dst_host_diff_srv_rate','dst_host_same_src_port_rate','dst_host_srv_diff_host_rate','dst_host_serror_rate','dst_host_srv_serror_rate','dst_host_rerror_rate','dst_host_srv_rerror_rate','outcome']

encode_numeric_zscore(df, 'duration')
encode_text_dummy(df, 'protocol_type')
encode_text_dummy(df, 'service')
encode_text_dummy(df, 'flag')
encode_numeric_zscore(df, 'src_bytes')
encode_numeric_zscore(df, 'dst_bytes')
encode_text_dummy(df, 'land')
encode_numeric_zscore(df, 'wrong_fragment')
encode_numeric_zscore(df, 'urgent')
encode_numeric_zscore(df, 'hot')
encode_numeric_zscore(df, 'num_failed_logins')
encode_text_dummy(df, 'logged_in')
encode_numeric_zscore(df, 'num_compromised')
encode_numeric_zscore(df, 'root_shell')
encode_numeric_zscore(df, 'su_attempted')
encode_numeric_zscore(df, 'num_root')
encode_numeric_zscore(df, 'num_file_creations')
encode_numeric_zscore(df, 'num_shells')
encode_numeric_zscore(df, 'num_access_files')
encode_numeric_zscore(df, 'num_outbound_cmds')
encode_text_dummy(df, 'is_host_login')
encode_text_dummy(df, 'is_guest_login')
encode_numeric_zscore(df, 'count')
encode_numeric_zscore(df, 'srv_count')
encode_numeric_zscore(df, 'serror_rate')
encode_numeric_zscore(df, 'srv_serror_rate')
encode_numeric_zscore(df, 'rerror_rate')
encode_numeric_zscore(df, 'srv_rerror_rate')
encode_numeric_zscore(df, 'same_srv_rate')
encode_numeric_zscore(df, 'diff_srv_rate')
encode_numeric_zscore(df, 'srv_diff_host_rate')
encode_numeric_zscore(df, 'dst_host_count')
encode_numeric_zscore(df, 'dst_host_srv_count')
encode_numeric_zscore(df, 'dst_host_same_srv_rate')
encode_numeric_zscore(df, 'dst_host_diff_srv_rate')
encode_numeric_zscore(df, 'dst_host_same_src_port_rate')
encode_numeric_zscore(df, 'dst_host_srv_diff_host_rate')
encode_numeric_zscore(df, 'dst_host_serror_rate')
encode_numeric_zscore(df, 'dst_host_srv_serror_rate')
encode_numeric_zscore(df, 'dst_host_rerror_rate')
encode_numeric_zscore(df, 'dst_host_srv_rerror_rate')
outcomes = encode_text_index(df, 'outcome')
num_classes = len(outcomes)
df.dropna(inplace=True,axis=1)

cDOS = [0,6,18,14,9,20]
cProbe = [5,15,10,17]
cU2R = [7,16,12,1]
cR2L = [3,8,2,19,13,4,21,22]

x, y = to_xy(df,'outcome')
print("loading model...")
model = load_model("trained_model.hdf5")

# Running random Test Case, takes random 10 inputs and predict the attack.
r=random.sample(range(0,len(x)),10)
for i in r:
    test_case = x[i]
    test_case = np.reshape(test_case,(-1,120))
    pred = model.predict(test_case)
    pred = np.argmax(pred, axis=1)
    c=bcolors.GREEN

    if int(pred) != 11:
        c=bcolors.RED

        pred = int(pred)
        if pred in cDOS:
            attackClass = "Dos"
        if pred in cProbe:
            attackClass = "Probes"
        if pred in cU2R:
            attackClass = "U2R"
        if pred in cR2L:
            attackClass = "R2L"
    else:
        attackClass = "normal"

    print(bcolors.WHITE+bcolors.BOLD+"\nTest case {} Input data >>".format(i)+bcolors.R+c+"\n{}\n{} ".format(rawcs.iloc[[i]],test_case)+bcolors.WHITE+"\nAttack type :"+bcolors.BOLD+c+" {}".format(outcomes[pred])+" attack Class: "+attackClass+"\n"+bcolors.R)
