### Escanear todo
### nmap -sU -p- -Pn 192.168.10.2
### Un unico puertos o intervalos
### nmap -sU -p 80 -T4 -Pn 192.168.10.2
### Con intervalor de envio
### nmap -sU -p 80-85 -T4 -Pn --scan-delay 3 192.168.10.2
### Puertos mas comunes
### nmap -sU -p 80-85 -Pn -F 192.168.10.2
### Tamaño payload
### nmap -sU -p 80-85 -Pn --data-length 64 192.168.10.2


from scapy.all import sniff, IP, UDP, ICMP, scapy
import time
import cachetools
import multiprocessing
import sys
import re



def waitresponsefunction(wait4response, lock):
    while 1:
        lock.acquire()
        to_remove = []
        if wait4response:
            for key in wait4response.keys():
                unix_time = time.time() 
                value = wait4response[key]
                if float(value['expire']) < unix_time:
                    print(len(wait4response))
                    ip_ori = key.split("||")[0]
                    file_name = "dataset_test.txt"
                    with open(file_name, "a") as file:
                        file.write(ip_ori + ',' + value['line'] + "\n") 
                    #wait4response.pop(key)
                    to_remove.append(key)
                    break
        for key in to_remove:
            wait4response.pop(key)       
        lock.release()




def process_packet(stateTable,wait4response,lock,maxWaitResponse,ip_addres_dst,packet):

    
    if IP in packet:
        ip_pkt = packet[IP]
        ip_id = ip_pkt.id
        packet_size = len(packet)
        current_time = packet.time

        
        if UDP in packet and len(packet.layers()) >= 4 and str(ip_pkt.dst) == ip_addres_dst:
            udp_pkt = packet[UDP]
            src_port = udp_pkt.sport
            dest_port = udp_pkt.dport
            name = str(ip_pkt.src) + '||' + str(src_port) + '||' + str(ip_pkt.dst) + '||' + 'UDP'
            name2waiting = str(ip_pkt.src) + '||' + str(src_port) + '||' + str(ip_pkt.dst) + '||' + str(dest_port) + '||' + 'UDP' + '||' + str(ip_id)

            if name in stateTable:
                
                ##Increase of connection ##Lock????
                valor = stateTable.pop(name)
                stateTable[name] = valor
                

                unix_time = time.time()  
                ##State of last package udp from that IP
                last_time_pkg = stateTable[name]['last']
                last_port = stateTable[name]['last_port']
                first_connection = stateTable[name]['first_connection']
                first_connection_total = stateTable[name]['first_connection_total']
                ##Calculate values of current package
                multiple_ports = 1 if last_port != dest_port else 0
                duration = str(current_time-first_connection) if last_port != dest_port else str(0)

                ##update value of connection
                first_connection_aux = current_time if last_port != dest_port else first_connection
                stateTable[name]['first_connection'] = first_connection_aux
                stateTable[name]['last'] = current_time
                stateTable[name]['last_port'] = dest_port
                ##Add to waiting
                line_to_csv = str(packet_size) + ',0,' + duration + ',' + str(current_time-first_connection_total)  + ',' + str(current_time-last_time_pkg) + ',0,' + str(multiple_ports) 
                wait4response[name2waiting] = {'line': line_to_csv, 'expire': unix_time + maxWaitResponse}
                

            else:

                unix_time = time.time()
                stateTable[name] = {'last':current_time, 'last_port':dest_port, 'first_connection' : current_time, 'first_connection_total' : current_time }

                line_to_csv = str(packet_size) + ',0,0,0,0,0,0'
                wait4response[name2waiting] = {'line': line_to_csv, 'expire': unix_time + maxWaitResponse}
    
   

        if ICMP in packet:
            icmp_pkt = packet[ICMP]
            icmp_type = icmp_pkt.type
            icmp_code = icmp_pkt.code
            if icmp_code == 3 and icmp_type == 3:
                if 'IP in ICMP' in icmp_pkt:
                    src_ip = icmp_pkt['IP in ICMP'].src
                    dest_ip = icmp_pkt['IP in ICMP'].dst
                    inside_id_ip = icmp_pkt['IP in ICMP'].id
                    if 'UDP in ICMP' in icmp_pkt and str(dest_ip) == ip_addres_dst:
                        src_port = icmp_pkt['UDP in ICMP'].sport
                        dest_port = icmp_pkt['UDP in ICMP'].dport
                        name2waiting = str(src_ip) + '||' + str(src_port) + '||' + str(dest_ip) + '||' + str(dest_port) + '||' + 'UDP' + '||' + str(inside_id_ip)
                        lock.acquire()
                        if name2waiting in wait4response:
                            auxLine = wait4response[name2waiting]['line']
                            auxSplit = auxLine.split(",")
                            auxSplit[1] = str(packet_size)
                            auxSplit[5] = str(1)
                            auxFinal = ','.join(auxSplit)
                            file_name = "dataset_test.txt"
                            with open(file_name, "a") as file:
                                file.write(str(src_ip) + ',' + auxFinal + "\n") 
                            wait4response.pop(name2waiting)
                        lock.release()


p_ipv4 = r'^(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$'


def is_double(string):
    try:
        float(string)
        return True
    except ValueError:
        return False

if __name__ == '__main__':

    arguments = sys.argv
    maxSizeCache = 0
    ttlCache = 0
    maxWaitResponse = 0
    ip_addres_dst = ''


    if len(arguments)==5 and arguments[1].isdigit() and arguments[2].isdigit() and is_double(arguments[3]):
        maxSizeCache = int(arguments[1])
        ttlCache = int(arguments[2])
        maxWaitResponse = float(arguments[3])
        if re.match(p_ipv4,arguments[4]):
            ip_addres_dst = arguments[4]
        else:
            print("Invalid IP")
            sys.exit()
    else:
        print('Invalid arguments')
        sys.exit()
    

    stateTable = cachetools.TTLCache(maxsize=maxSizeCache, ttl=ttlCache)
    admin = multiprocessing.Manager()
    lock = multiprocessing.Lock()
    wait4response = admin.dict({})

    process = multiprocessing.Process(target=waitresponsefunction, args=(wait4response,lock))
    process.start()


    sniff(filter="udp or icmp", prn=lambda packet: process_packet(stateTable,wait4response,lock,maxWaitResponse,ip_addres_dst,packet))