import sys
import os
import json
import subprocess
import pprint
import time
import datetime
import pandas as pd
from datetime import datetime
from bat import bro_log_reader
from bat.log_to_dataframe import LogToDataFrame
from pyod.models.pca import PCA

# TODO s: 
# - Seach and resolve others TODO s in whole code.
# - Analyze if internal and network IPs can be read from another file (like README.md or come conf) or can be hardcoded as 192.168.X.X
# - Create a label hiarerchy class like: father label class with name, characteristics and points, then label childrens for levels and customs.
# - Clean code (prints, wrong comments, main, help, params, hardcoded variables)

# Bro-Dataframe Rows Format:
# [0] uid                    CXvywd555Le9IImG9
# [1] id.orig_h                  192.168.1.197
# [2] id.orig_p                          51656
# [3] id.resp_h                 185.244.25.108
# [4] id.resp_p                             80
# [5] proto                                tcp
# [6] service                             http
# [7] duration          0 days 00:00:01.954395
# [8] orig_bytes                            86
# [9] resp_bytes                         51041
# [10] conn_state                            SF
# [11] local_orig                           NaN
# [12] local_resp                           NaN
# [13] missed_bytes                           0
# [14] history                        ShADadtfF
# [15] orig_pkts                             80
# [16] orig_ip_bytes                       5492
# [17] resp_pkts                             78
# [18] resp_ip_bytes                     106154
# [19] tunnel_parents                   (empty)
# [20] MultiLabel                          None

pwd = os.popen('pwd').read()[:-1]
csvfile = "labeled_dataframe"
logfile = str(pwd) + "/labeled_dataframe_assigner.log"
date = os.popen('date').read()[:-1]
errorlog =  "No errors."
success = ""

# Level Labeling Function
# 1st Level (direction): From, To.
# 2nd Level (host): Background, Bening, InfectedDevice.
# 3rd Level (Transport Layer Protocol): UDP, TCP.
# 4th Level (Connection State): Established, Attempts. 
# 5th Level (Application Layer Protocol): HTTP, DNS, FTP, Telnet, SSH, HTTP/TLS
# 6th Level (Encoded/Encrypted): BinaryProtocol (this is for encrypted), Encoded, ASCII (or empty), Encrypted (you should be sure)
# 7th Level (Extra Comments): Upload, Download, FileDownload, HTTP-Ad, SPAM,  Proxy, CC, DDoS, HeartBeat, ConnectivityCheck. (can be empty) 
def level_labeling(dataframe_row):
    row = dataframe_row
    labeled_row1 = first_level_labeling(row)
    labeled_row2 = second_level_labeling(labeled_row1)
    labeled_row3 = third_level_labeling(labeled_row2)
    labeled_row4 = fourth_level_labeling(labeled_row3)
    labeled_row5 = fifth_level_labeling(labeled_row4)
    labeled_row6 = sixth_level_labeling(labeled_row5)
    return labeled_row6

# 1st Level (direction): 
# FromInfectedDevice: from local infected device. To automate this we need the local infected device IP and see if it its the source IP in the connection flow. 
# ToInfectedDevice: to local infected device. To automate this we need the local infected device IP and see if it its the destination IP in the connection flow.
def first_level_labeling(dataframe_row):
    #print("first_level_labeling")
    row = dataframe_row
    # Getting the network and host IP from Bro log
    readed_ip = row[1][1]
    network_third_octet_ip = "" 
    internal_network = "192.168." + str(network_third_octet_ip)
    if internal_network in readed_ip:
        ip = ""
        ip = ip.join(readed_ip.split(".")[2:-1])
        ip += ip.join(".")
        ip += ip.join(readed_ip.split(".")[3:])
        host_last_octet_ip = ip
    else:
        host_last_octet_ip = "0.1"
    internal_ip = str(internal_network) + str(host_last_octet_ip)
    
    # From InfectedDevice
    # [1] id.orig_h must be -> internal
    # [3] id.resp_h must be -> external
    if (row[1][1] == internal_ip) and (row[1][3] != internal_network):
        row[1][int(columns)] += "-FromInfectedDevice"
        row[1][int(columns)+1] += "FromInfectedDevice"

    # To InfectedDevice
    # [1] id.orig_h must be -> external
    # [3] id.resp_h must be -> internal
    if (row[1][3] == internal_ip) and row[1][3] (row[1][1] != internal_network):
        row[1][int(columns)] += "-ToInfectedDevice"
        row[1][int(columns)+1] += "ToInfectedDevice"

    return row

# 2nd Level (type):
# Malicious (Automated): All flows in malware captures puts a malicious label in the Label column.   
# Bening (Bening): benning flows in malware traffic captures and all flows in bening captures
# InfectedDevice puts a bening label in the Label column.
def second_level_labeling(dataframe_row):
    #print("second_level_labeling")
    # This has te be manual
    row = dataframe_row
    return row

# 3rd Level (Transport Layer Protocol): UDP, TCP.
def third_level_labeling(dataframe_row):
    #print("third_level_labeling")
    # Could be removed to avoid information redundancy
    # [5] proto tcp/udp
    row = dataframe_row
    if (row[1][5] == "tcp"):
        row[1][int(columns)] += "-TCP"
        row[1][int(columns)+1] += "TCP"
    
    if (row[1][5] == "udp"):
        row[1][int(columns)] += "-UDP"
        row[1][int(columns)+1] += "UDP"
   
    return row

# 4th Level (Connection State): Established, Attempts. 
def fourth_level_labeling(dataframe_row):
    #print("fourth_level_labeling")
    # TODO: Analyze how to label this states maybe with conn_state (see what the letters means)
    # [10] conn_state   SF
    # TODO: Analyze how to label this states maybe with duration, talk with Frenky and Pau
    # [7] duration  0 days 00:00:01.954395
    row = dataframe_row
    
    established_states =  ["S1", "SF", "S2", "S3", "RSTO", "RSTR"]
    attempts_states = ["S0", "REJ", "RSTOS0", "RSTRH", "SH", "SHR", "OTH"]

    try:
        if str(row[1][10]) in established_states:
            #print("FOUND ESTABLISHED")
            row[1][int(columns)] += "-Established"
            row[1][int(columns)+1] += "Established"
    
        if str(row[1][10]) in attempts_states:
            #print("FOUND ATTEMP")
            row[1][int(columns)] += "-Attempt"
            row[1][int(columns)+1] += "Attempt"
    
        if str(row[1][7]) != "NaT":
            time_row = str(row[1][7]).split("days")[1].strip()
            date_struct = datetime.strptime(str(time_row), '%H:%M:%S.%f')
            ms = date_struct.microsecond
        
            # 500000 Microseconds = half second
            if (ms > 500000):
                row[1][int(columns)] += "-DurationHalfSec"
                row[1][int(columns)+1] += "DurationHalfSec" 
        else:
            row[1][int(columns)] += "-NoDurationInfo"
            row[1][int(columns)+1] += "NoDurationInfo"
    except:
        #print(str(row[1]))
        #print(str(row[1][7]))
        time.sleep(5)

    return row

# 5th Level (Application Layer Protocol): HTTP, DNS, FTP, Telnet, SSH, HTTP/TLS, IRC
def fifth_level_labeling(dataframe_row):
    #print("fhttps://miro.com/app/board/o9J_kzMPiWE=/ifth_level_labeling")
    # [6] service   http
    row = dataframe_row
    if (row[1][6] == "http"):
        #print("FOUND HTTP")
        row[1][int(columns)] += "-HTTP"
        row[1][int(columns)+1] += "HTTP"

    row = dataframe_row
    if (row[1][6] == "ssh"):
        #print("FOUND SSH")
        row[1][int(columns)] += "-SSH"
        row[1][int(columns)+1] += "SSH"

    row = dataframe_row
    if (row[1][6] == "irc"):
        #print("FOUND IRC")
        row[1][int(columns)] += "-IRC"
        row[1][int(columns)+1] += "IRC"

    row = dataframe_row
    if (row[1][6] == "telnet"):
        #print("FOUND TELNET")
        row[1][int(columns)] += "-Telnet"
        row[1][int(columns)+1] += "Telnet"

    row = dataframe_row
    if (row[1][6] == "dns"):
        #print("FOUND DNS")
        row[1][int(columns)] += "-DNS"
        row[1][int(columns)+1] += "DNS"
    
    row = dataframe_row
    if (row[1][6] == "ftp"):
        #print("FOUND FTP")
        row[1][int(columns)] += "-FTP"
        row[1][int(columns)+1] += "FTP"

    # TODO: check how bro detects HTTP/TLS
    row = dataframe_row
    if (row[1][6] == "tls"):
        #print("FOUND HTTP/TLS")
        row[1][int(columns)] += "-HTTP/TLS"
        row[1][int(columns)+1] += "HTTP/TLS"

    return row

# Manual Labeling only
# 6th Level (Encoded/Encrypted): BinaryProtocol (this is for encrypted), Encoded, ASCII (or empty), Encrypted (you should be sure)
def sixth_level_labeling(dataframe_row):
    #MANUAL 6th Level (Encrypted): 
    
    #BinaryProtocol (this is for encrypted): if we see in the analysis that the payload is binary.
    #Encoded: if we see in the analysis that the payload is encoded with any encoding system.
    #ASCII (or empty): if we see in the analysis that the payload is in ASCII (this could be the default behaviour and not label feature).
    #Encrypted (you should be sure): Only if we are sure after a deep analysis that it is encrypted with some encryption system different than SSL and its not only encoded.

    return row

# Manual Labeling only
# 7th Level (Extra Comments): Upload, Download, FileDownload, HTTP-Ad, SPAM, Proxy, CC, DDoS, HeartBeat, ConnectivityCheck. (can be empty) 
def seventh_level_labeling(dataframe_row):
    #MANUAL 7th Level (Extra Comments): 
    
    #Upload: if we see in the analysis that the source is the local infected device and it is uploading something.
    #Download: if we see in the analysis that the source is the local infected device and it is downloading something.
    #FileDownload: The same as Download???? TODO
    #HTTP-Ad: DONT KNOW HOW TO DETECT THIS, EXAMPLE? TODO
    #SPAM: DONT KNOW HOW TO DETECT THIS, EXAMPLE?  TODO
    #Proxy: if we see in the analysis the the destination IP is then a source IP to a third host/server and back. NOT SURE
    #C&C: if we see in the analysis that our infected device has a flows to some IP in which the payload has a suspicious communicational behaviour and its seem to be sending commands, retrieving data and/or downloading files.
    #DDoS: if we see in the analysis that some destination IP has a lot of flows with few packets, same number of bytes from multiple different source IPs. We could see other flow/packets characteristics to detect different types of DDoS like SYN flag, Connection States, multiple source IPs and few destination IPs at the same ports, etc. Interesting source: Detecting and Mitigating DDOS Attacks
    #HeartBeat: if we see in the analysis that the suspected C&C IP has flows  with few bytes and few packets with the infected device. 
    #ConnectivityCheck: Same as HeartBeat? Or maybe the inverse? TODO
    #PartOfHorizontalPortScan: Is wee see in the analysis that there are many flows with source IP our infected device and different destination IPs at the same port but with a not established conn state/label 
    #Attack: Is wee see in the analysis that there are many flows with source IP our infected device and different destination IPs at the same port but with a established conn state/label 
    #Empty: when none of the other characteristics on all the levels set any label.

    return row


# Custom Labeling Function (testing purposes)
def custom_labeling(dataframe_row, configfile_ok):
    #print("custom_labeling")
    row = dataframe_row
    labeled_row = port_labeling(row)
    labeled_row = service_labeling(labeled_row)
    labeled_row = payload_labeling(labeled_row)

    if configfile_ok == 1:
        try:
            with open('outfile_labeling.json') as json_file:
                data = json.load(json_file)
                labeled_row = configfile_labeling(labeled_row, data)
        except Exception as e:
            errorlog = str(e)
            print("Exception in reading config file. Turning off configfile labeling " + str(e))
            configfile_ok = 0
    
    t = (labeled_row,configfile_ok)
    return t

# This function put labels getting the information from and out config file. This is for new findings in capture analysis.
# Out config file format: #{"data":[{"numner":4, "comparator":"4441", "function":"eq","label":"Mirai","type":"Malicious"},{"numner":4, "comparator":"53413", "function":"eq","label":"Okane","type":"Malicious"}]}
def configfile_labeling(dataframe_row, data):
    #print("port_labeling")
    row = dataframe_row
    
    try:

        for bl in data['data']:
            brolog_feature_number = bl['number'] #bro log feature number: 4 for port
            comparator = bl['comparator'] #compare with: 4441 for Mirai port
            comp_function = bl['function'] #how to compare: ==
            label = bl['label'] #label: Mirai
            flow_type = bl['type'] #type: Malicious or Benign

            try:
                comparator = int(comparator)
            except:
                comparator = str(comparator)

            if str(comp_function) == "eq":    
                if row[1][int(brolog_feature_number)] == comparator:
                    row[1][int(columns)] += "-" + str(label)
                    row[1][int(columns)+1] = str(label) 
                    row[1][int(columns)+2] = str(flow_type) 
        
            if str(comp_function) == "lt":    
                if row[1][int(brolog_feature_number)] < comparator:
                    row[1][int(columns)] += "-" + str(label)
                    row[1][int(columns)+1] = str(label) 
                    row[1][int(columns)+2] = str(flow_type) 
        
            if str(comp_function) == "gt":    
                if row[1][int(brolog_feature_number)] > comparator:
                    row[1][int(columns)] += "-" + str(label)
                    row[1][int(columns)+1] = str(label) 
                    row[1][int(columns)+2] = str(flow_type) 
        
    except Exception as e:
        errorlog = str(e)
        print("Exception in reading json data at configfile_labeling function" + str(e))
        raise(e)

    return row

# Port labeling can be updated as manul analysis grows.
def port_labeling(dataframe_row):
    #print("port_labeling")
    row = dataframe_row
    
    # Known Mirai Port 4441
    if row[1][4] == 4441:
        row[1][int(columns)] += "-MiraiPort"
        row[1][int(columns)+1] = "MiraiPort" 
        row[1][int(columns)+2] = "Malicious" 


    # Known Router Attack Mirai Variant Port 53413        print(a + b - c)

    if row[1][4] == 53413:
        row[1][int(columns)] += "-MiraiPort"
        row[1][int(columns)+1] = "MiraiPor" 
        row[1][int(columns)+2] = "Malicious" 
           
    
    # Known Router Attack Satori (Mirai Variant) Port 37215
    if row[1][4] == 37215:
        row[1][int(columns)] += "-SatoriPort"
        row[1][int(columns)+1] = "SatoriPort"    
        row[1][int(columns)+2] = "Malicious" 

    # Known Router Attack Satori (Mirai Variant) Port 52869
    if row[1][4] == 52869:
        row[1][int(columns)] += "-SatoriPort"
        row[1][int(columns)+1] = "SatoriPort"  
        row[1][int(columns)+2] = "Malicious" 

    # Known Steam Attack Okane (Mirai Variant) Port 27015
    if row[1][4] == 27015:
        row[1][int(columns)] += "-OkanePort-RouterAttack"
        row[1][int(columns)+1] = "OkanePort"  
        row[1][int(columns)+2] = "Malicious" 

    return row

def service_labeling(dataframe_row):
    #print("service_labeling")
    row = dataframe_row

    # SSH Service (need to know how to distinguish from LAN and WAN SSH connections and directions)
    if str(row[1][6]) == "ssh":
        row[1][int(columns)] += "-SSHExternal"
        row[1][int(columns)+1] = "SSHExternal" 

    # DNS Service (testing purposes)
    if str(row[1][6]) == "dns":
        row[1][int(columns)] += "-DNSQuery"
        row[1][int(columns)+1] = "DNSQuery" 

    # IRC Service
    if str(row[1][6]) == "IRC":
        row[1][int(columns)] += "-IRCChannel"
        row[1][int(columns)+1] = "IRCChannel" 

    # Telnet Service
    if str(row[1][6]) == "Telnet":
        row[1][int(columns)] += "-TelnetControl"
        row[1][int(columns)+1] = "TelnetControl" 

    return row

def payload_labeling(dataframe_row):
    #print("payload_labeling")
    row = dataframe_row

    # Big payload data
    if str(row[1][9]) != "nan":
        if int(row[1][9]) > 5000000:
            row[1][int(columns)] += "-BigPayload"
            row[1][int(columns)+1] = "BigPayload"  

    # No payload data
    if str(row[1][9]) != "nan":
        if int(row[1][9]) < 1:
            row[1][int(columns)] += "-NoPayload"
            row[1][int(columns)+1] = "NoPayload" 

    # TODO: See how to retrieve payload data and get the hexadecimal to compare with known malicious hexa codes.

    return row


if __name__ == '__main__':
    x = datetime.now()
    print("Starting script at {}".format(x))
    
    configfile_ok = 1
    pp = pprint.PrettyPrinter(indent=4)
    new_list = []

    brologfile = sys.argv[1]
    print('Analizing Bro File: ', str(brologfile))
    csvfile = str(pwd) + "/" + str(brologfile).split("/")[0] + "/" + str(brologfile).split("/")[-1] + "_" + str(csvfile) + ".csv"
    print("Results at: " + str(csvfile))    
    print("Logs at: " + str(logfile))
    
    # Create a Pandas dataframe from a Bro log
    log_to_df = LogToDataFrame()
    bro_df = log_to_df.create_dataframe(brologfile)
    
    # Insert new labels columns at the end
    columns=len(bro_df.columns)
    bro_df.insert(int(columns), "MultiLabel", "", True) 
    bro_df.insert(int(columns)+1, "UniLabel", "", True)
    bro_df.insert(int(columns)+2, "Label", "", True) 
    
    # Loop the generated dataframe rows to search for key analysis patterns (such as known malicious ports)
    iterrows = bro_df.iterrows()    
    information=0
    leniterrows=os.popen('wc -l ' + str(brologfile) + '| cut -d " " -f1 ').read()
    print("Analyzing a total of {} this may take several minutes or hours . . .".format(leniterrows))

    try:
        for row in iterrows:
            labeled_row = custom_labeling(row, configfile_ok)
            configfile_ok = labeled_row[1]
            labeled_row = level_labeling(labeled_row[0])
            new_list.append(labeled_row[1])
            information = information + 1
            if (information % 10000) == 0:
                print("Analyzing {} rows of {} at {}".format(information, leniterrows, datetime.now()))

        new_df = pd.DataFrame(new_list, columns = list(bro_df.columns))
        #csvfile = csvfile + str(brologfile)
        new_df.to_csv(csvfile)
        sucess="Results written in " + str(csvfile)

    except Exception as e:
        errorlog = str(e)
        sucess = "Error has ocurred."
        raise(e)


log_file= open(str(logfile),"a+")
log_file.write("----- Logs from last execution of labeled_dataframe_assigner.py -----\n")
log_file.write("Date: " + str(date) + "\n")
log_file.write("Errors: " + str(errorlog) + "\n")
log_file.write("Sucess: " + str(sucess) + "\n")
log_file.write("---------- Logs Ends ----------\n")
