# -*- coding: utf-8 -*-
"""GradioInterface_v2.ipynb

Automatically generated by Colaboratory.
"""

# Commented out IPython magic to ensure Python compatibility.
# # Capture to supress the download ouput
# %%capture
# !pip install gradio
# !pip install pandas
# !pip install transformers
# !pip install parsezeeklogs
# !pip install elasticsearch

# Define imports for model use
import torch
from transformers import pipeline
from parsezeeklogs import ParseZeekLogs
from transformers import BertTokenizer
import gradio as gr
import pandas as pd

# Define model
pipe = pipeline(model="19kmunz/IoT-23-BERT-Network-Logs-Classification", tokenizer=BertTokenizer.from_pretrained("bert-base-cased"))

# Define string constants
LOG = "conn.log Output"
HEADER_TABLE = "Headers Table"
SENTENCES = "Sentences"
OUT = "out"
INPUT_TYPES = [LOG, HEADER_TABLE, SENTENCES]
STEPS = [HEADER_TABLE, SENTENCES]
HEADERS=['id.resp_p', 'proto', 'conn_state', 'orig_pkts', 'orig_ip_bytes', 'resp_ip_bytes']

# Define sentence-ization functions
# Dictionary of feature names to use in the make sentence function
feature_names = {'id.resp_p':'response port',
                 'proto':'transport protocol',
                 'orig_pkts':'number of packets sent by the origin',
                 'conn_state':'connection state',
                 'orig_ip_bytes':'number of IP level bytes sent by the originator',
                 'resp_ip_bytes':'number of IP level bytes sent by the responder'}

# Function to make sentences out of the data
def make_sentence(row):
  sentences = {}
  for feature in row.keys():
    if feature == 'label' or feature == "#":
      sentences[feature] = row[feature]
    else:
      sentences[feature] = feature_names[feature] + " is " + str(row[feature]) + "."
  return sentences

# Take all sentence observations and make them into paragraph inputs
def make_paragraphs(ser):
  paragraphs_list = []
  for index,obs in ser.items():
    new_para = obs['id.resp_p'] + " " + obs['proto'] + " " + obs['conn_state'] + " " + obs['orig_pkts'] + " " + obs['orig_ip_bytes'] + " " + obs['resp_ip_bytes']
    paragraphs_list.append(new_para)
  return pd.Series(paragraphs_list, name="Sentences", index=ser.index).to_frame()

# Define prediction Functions For Different Settings
def predictFromSentences(sentenceTable):
  output = pipe(sentenceTable[SENTENCES].tolist()) # This does the prediction!
  return { OUT: pd.DataFrame({"Output": ["Malicious" if pred['label'] == "LABEL_0" else "Benign" for pred in output] }) }

def predictFromHeaderTable(headerTable):
  sentences = headerTable.apply(make_sentence, axis=1);
  paragraphs = make_paragraphs(sentences)
  return {
      SENTENCES: paragraphs,
      OUT: predictFromSentences(paragraphs)[OUT]
  }

def predictFromFileUpload(fileUpload):
  if(fileUpload is None):
      raise gr.Error("No file uploaded")
  fileType = fileUpload.split('.')[-1]
  if(fileType == 'csv'):
    dataFrame = pd.read_csv(fileUpload, usecols=HEADERS)
  elif(fileType == 'log' or fileType == 'labeled'):
    with open('out.csv',"w") as outfile:
      for log_record in ParseZeekLogs(fileUpload, output_format="csv", safe_headers=False, fields=HEADERS):
          if log_record is not None:
              outfile.write(log_record + "\n")
    dataFrame = pd.read_csv('out.csv', names=HEADERS)
  result = predictFromHeaderTable(dataFrame)
  toReturn =  {
      HEADER_TABLE: dataFrame,
      SENTENCES: result[SENTENCES],
      OUT: result[OUT]
  }
  return toReturn

def makeIndexColumn(allInputs):
  def _makeIndexColumnFor(column):
    theseHeaders = allInputs[column].columns
    newHeaders = ['#', *theseHeaders]
    allInputs[column]['#'] = allInputs[column].index
    allInputs[column] = allInputs[column][newHeaders]

  if(SENTENCES in allInputs):
    _makeIndexColumnFor(SENTENCES)
  if(HEADER_TABLE in allInputs):
    _makeIndexColumnFor(HEADER_TABLE)
  if(OUT in allInputs):
    _makeIndexColumnFor(OUT)
  return allInputs

def predict(inputType, fileUpload, headerTable, sentenceTable, out):
  output = {};
  if(inputType == LOG):
    # Process File Upload
    output = makeIndexColumn(predictFromFileUpload(fileUpload))
    return [output[HEADER_TABLE], output[SENTENCES], output[OUT]]
  elif(inputType == HEADER_TABLE):
    # Process Header Table
    output = makeIndexColumn(predictFromHeaderTable(headerTable))
    return [headerTable, output[SENTENCES], output[OUT]]
  elif(inputType == SENTENCES):
    # Process From Sentences
    output = makeIndexColumn(predictFromSentences(sentenceTable))
    return [headerTable, sentenceTable, output[OUT]]

# Update UI
def updateInputOutputBlocks(inputType, steps):
  # Update visibility and Interactivity of Gradio Blocks based on Settings
  fileUpload = gr.File(
      visible=(True if inputType == LOG else False),
      interactive=(1 if inputType == LOG else 0)
  )
  headerTable = gr.Dataframe(
      visible=(True if (inputType == HEADER_TABLE or HEADER_TABLE in steps) else False),
      interactive=(1 if inputType == HEADER_TABLE else 0)
    )
  sentenceTable = gr.Dataframe(
      interactive=(1 if inputType == SENTENCES else 0),
      visible=(True if (inputType == SENTENCES or SENTENCES in steps) else False)
    )
  return fileUpload, headerTable, sentenceTable

# Create Gradio UI
with gr.Blocks() as app:
    gr.Markdown("""
    # Network Log Predictions
    Input log information below and click 'Run' to get predictions from our model!
    Access the settings at the bottom for different types of input and to see inbetween steps.
    """)
    # Inputs / Outputs
    fileUpload = gr.File(file_types=[".log", ".log.labeled", ".csv"], label="Zeek Log File", visible=False, file_count='single')
    headerTable = gr.Dataframe(row_count = (2, "dynamic"), col_count=(7,"fixed"), headers=['#', *HEADERS], label="Header Inputs", interactive=1)
    sentenceTable = gr.Dataframe(row_count = (2, "dynamic"), col_count=(2, "fixed"), headers=["#", "Sentence"], label="Sentences", interactive=0, visible=False)
    out = gr.Dataframe(row_count = (2, "dynamic"), col_count=(2, "fixed"), headers=['#', "Output"], label="Predictions", column_widths=["60px", "100%"])
    btn = gr.Button("Run")

    # Settings
    with gr.Accordion("Settings", open=False):
      inputType = gr.Radio(INPUT_TYPES, value="Headers Table", label="Input")
      steps = gr.CheckboxGroup(STEPS, label="Display Intermediary Steps")
      inputType.change(
        fn=updateInputOutputBlocks,
        inputs=[inputType, steps],
        outputs=[fileUpload, headerTable, sentenceTable]
      )
      steps.change(
        fn=updateInputOutputBlocks,
        inputs=[inputType, steps],
        outputs=[fileUpload, headerTable, sentenceTable]
      )
    # Assign Callback
    btn.click(
        fn=predict,
        inputs=[inputType, fileUpload, headerTable, sentenceTable, out],
        outputs=[headerTable, sentenceTable, out]
        )

app.launch()