from scapy.all import *

def TrackState(Packet, Key, Ephemeral):

    #Create key entry in dict if it doesn't exist
    if not Key in ConnectionDict:
        ConnectionDict[Key] = {}

    if Packet.sprintf("%TCP.flags%") == "A":
        if Ephemeral in ConnectionDict[Key]: #Does ephemeral entry exist for this key?
            if ConnectionDict[Key][Ephemeral] == "FA" or ConnectionDict[Key][Ephemeral] == "FPA": #FA-ACK, FPA-ACK
                ConnectionDict [Key].pop(Ephemeral) #Delete Ephemeral port entry if previous message was FA and this one was ACK                    
            else: #A - not FA, or FPA
                ConnectionDict [Key][Ephemeral] = Packet.sprintf("%TCP.flags%") #ACK seen, entry for key exists, revising entry                    
        else:
            ConnectionDict [Key][Ephemeral] = Packet.sprintf("%TCP.flags%") #ACK seen without previous entry for key
    elif Packet.sprintf("%TCP.flags%") == "R": #Reset packet Seen
        if Ephemeral in ConnectionDict[Key]:
            ConnectionDict [Key].pop(Ephemeral)  #Delete Ephemeral port entry if Reset flag is given                
    else: #Every flag except A and R
        ConnectionDict [Key][Ephemeral] = Packet.sprintf("%TCP.flags%")

    #Delete key if no other entries exist
    if len(ConnectionDict[Key]) == 0:
        ConnectionDict.pop(Key)

    return


def MonitorTraffic(Packet):

    #Create a consistent key, ephemeral, and FromServer (Used for termination) variables
    if Packet.haslayer(TCP):
        if Packet[TCP].dport == 21:
            Key = "{}-{}".format(Packet[IP].src,Packet[IP].dst)
            Ephemeral = Packet[TCP].sport
            FromServer = False
        elif Packet[TCP].sport == 21:
            Key = "{}-{}".format(Packet[IP].dst,Packet[IP].src)
            Ephemeral = Packet[TCP].dport
            FromServer = True
        else: #Not FTP Traffic
            return

        #If packet comes back in from a session that was suppose to be terminated, terminate it again.
        if Key in TerminateDict:
            if Ephemeral in TerminateDict[Key]:
                if not (FromServer and Packet.sprintf("%TCP.flags%") == "R"):
                    TerminateSession(Packet)
                    return

        if 'ConnectionDict' in globals(): #If state tracking is enabled
            TrackState(Packet, Key, Ephemeral)

        if Packet.haslayer(Raw): #Parse Commands
            if ((Packet[Raw].load).find(b"230")) != -1: #Termination triggered on successful login
                TerminateSession(Packet) #Initial Termination
                if not Key in TerminateDict:
                    TerminateDict[Key] = {}
                if not Ephemeral in TerminateDict[Key]:
                    TerminateDict [Key][Ephemeral] = b"TERM"

        #Print Contents of variables, Packet Flags and Commands
        if debugging:
            if Packet.haslayer(Raw):
                print(Key, Ephemeral, Packet.sprintf("%TCP.flags%"), Packet[Raw].load.strip())
            else: 
                print(Key, Ephemeral, Packet.sprintf("%TCP.flags%"))


def TerminateSession(Packet):
    ack = Packet[TCP].seq+1
    seq = Packet[TCP].ack
    resetpacket = Ether(src=Packet[Ether].dst, dst=Packet[Ether].src)/IP(src=Packet[IP].dst, dst=Packet[IP].src)/TCP(sport=Packet[TCP].dport, dport=Packet[TCP].sport, seq=seq, ack=ack, flags="R", options=Packet[TCP].options)
    sendp(resetpacket, verbose=False)
    if debugging:
        print("Terminating Session")

def main():

    StateTracking = False #Track the TCP Flag state of the Packet and Session

    if StateTracking:
        global ConnectionDict
        ConnectionDict = {}

    global TerminateDict
    TerminateDict = {}    

    global debugging
    debugging = True

    if not debugging:
        print ("Monitoring for Traffic....Hit Break to Quit")

    try:
        sniff(filter="tcp port 21", prn=MonitorTraffic)
    except Exception as e:
        print(e)
    #finally:
    #    print(ConnectionDict)
    #    print(TerminateDict)



if __name__ == "__main__":
    main()


