from scapy.all import *

def Full_Duplex_Sessions(p):
    ##Source : Mark Baggett - https://gist.github.com/MarkBaggett/d8933453f431c111169158ce7f4e2222#file-scapy_helper-py
    sess = "Other"
    if 'Ether' in p:
        if 'IP' in p:
            if 'TCP' in p:
                sess = str(sorted(["TCP", p[IP].src, p[TCP].sport, p[IP].dst, p[TCP].dport],key=str))
            elif 'UDP' in p:
                sess = str(sorted(["UDP", p[IP].src, p[UDP].sport, p[IP].dst, p[UDP].dport] ,key=str))
            elif 'ICMP' in p:
                sess = str(sorted(["ICMP", p[IP].src, p[IP].dst, p[ICMP].code, p[ICMP].type, p[ICMP].id] ,key=str)) 
            else:
                sess = str(sorted(["IP", p[IP].src, p[IP].dst, p[IP].proto] ,key=str)) 
        elif 'ARP' in p:
            sess = str(sorted(["ARP", p[ARP].psrc, p[ARP].pdst],key=str)) 
        else:
            sess = p.sprintf("Ethernet type=%04xr,Ether.type%")

    return sess




def TCPSessionFilter(sessions):
#Pulls out only TCP sessions from raw session list

    Keys=list(sessions.keys())
    Values=list(sessions.values())

    TCPSessions = {}

    for (KeyEntry, ValueEntry) in zip(Keys,Values):        
        Key = KeyEntry.replace("[","").replace("]","").replace("'","").replace(", ","-").split("-")
        
        if len(Key) == 5: #Does the session entry have 5 values
            if Key[4] == "TCP": #Is the fifth one "TCP"?
                TCPSessions [KeyEntry] = ValueEntry
    
    return TCPSessions

def FTPSessionFilter(sessions):
#Pulls out only FTP sessions from TCP session list

    Keys=list(sessions.keys())
    Values=list(sessions.values())

    FTPSessions = {}
    SaveFTPSession = False

    for Kindex,Vindex in zip(Keys, Values): 
        for Yindex in Vindex:
            if (Yindex.haslayer(TCP)): #A bit redundant but why not?
                if Yindex[TCP].dport == 21 or Yindex[TCP].sport == 21: #FTP Packets Found
                    SaveFTPSession = True
        
        if (SaveFTPSession): #Save session with FTP packets
            FTPSessions [Kindex] = Vindex
            SaveFTPSession = False #Reset 
    

    return FTPSessions

def FTPBruteAnalysis(FTPSessions, plogins=False, psummary=False, savefile=False, trackattacks=False):
#Analyizes an FTP stream list for login failures and prints or saves them

    Keys = list(FTPSessions.keys())
    Values = list(FTPSessions.values())


    User = b''
    Password = b''
    FTPsuccess = 0
    FTPfail = 0
    FTPfailpackets = {}
    if trackattacks:
        FTPattacks = {}

    for Kindex,Vindex in zip(Keys, Values): #Sifting through each session
        FTPsave = False #Reset per session
        for Yindex in Vindex: #Sifting through each packet in the session
            if (Yindex.haslayer(Raw)): #If the packet doesn't have a Raw layer, ignore          
                if ((Yindex[Raw].load).find(b"USER")) != -1: #Found USER in pkt
                    User = Yindex[Raw].load.strip().split()[1]
                elif ((Yindex[Raw].load).find(b"PASS")) != -1: #Found PASS in pkt
                    Password = Yindex[Raw].load.strip().split()[1]
                elif ((Yindex[Raw].load).find(b"230 Login successful")) != -1:
                    if plogins:
                        print("Found LOGIN SUCCESS in Session")
                        print("Client :", Yindex[IP].dst ,"Server :", Yindex[IP].src)
                        print ("Username :", User, " Password :", Password)
                        User = b''
                        Password = b''                        
                    FTPsuccess += 1
                elif ((Yindex[Raw].load).find(b"530 Login incorrect")) != -1:
                    if plogins:
                        print("Found LOGIN FAILURE in Session")
                        print("Client :", Yindex[IP].dst ,"Server :", Yindex[IP].src)
                        print ("Username :", User, " Password :", Password)
                        User = b''
                        Password = b''                    
                    if trackattacks:
                        AttackKey = "{}->{}".format(Yindex[IP].dst,Yindex[IP].src)

                        if AttackKey in FTPattacks.keys():
                            FTPattacks [AttackKey] = FTPattacks [AttackKey] + 1
                        else:
                            FTPattacks [AttackKey] = 1


                    FTPfail += 1
                    FTPsave = True
            
        if (FTPsave):
            FTPfailpackets [Kindex] = Vindex


    if psummary:
        print ("\nFTP Success Count :", FTPsuccess)
        print ("FTP Failure Count :", FTPfail)
    
    if trackattacks:
        print ("\n")
        for key, value in FTPattacks.items():
            print(key, ' : ', value)

    if savefile:
        SaveSessionsToFile(FTPfailpackets)


def SSHSessionFilter(sessions):
#Pulls out only SSH sessions from TCP session list

    Keys=list(sessions.keys())
    Values=list(sessions.values())

    SSHSessions = {}
    SaveSSHSession = False

    for Kindex,Vindex in zip(Keys, Values): 
        for Yindex in Vindex:
            if (Yindex.haslayer(TCP)): #A bit redundant but why not?
                if Yindex[TCP].dport == 22 or Yindex[TCP].sport == 22: #SSH Packets Found
                    SaveSSHSession = True
        
        if (SaveSSHSession): #Save session with SSH packets
            SSHSessions [Kindex] = Vindex
            SaveSSHSession = False #Reset 
    

    return SSHSessions


def SSHBruteAnalysis(SSHSessions, threshold = 7500, psummary=False, trackattacks = False):
#Analyizes an SSH stream list for potential login failures

    Keys = list(SSHSessions.keys())
    Values = list(SSHSessions.values())


    SSHsuccess = 0
    SSHfail = 0
    SSHFailSessions = {}
    if trackattacks:    
        SSHattacks = {}

    for Kindex,Vindex in zip(Keys, Values): #Sifting through each session

        SessionLength = 0
        for Yindex in Vindex: #Sifting through each packet in the session
            SessionLength += (Yindex[IP].len + 14)

        if SessionLength < threshold:
            SSHFailSessions [Kindex] = Vindex
            SSHfail += 1
        else:
            SSHsuccess += 1

        if trackattacks:
            AttackKey = "{}->{}".format(Yindex[IP].src,Yindex[IP].dst)
            if AttackKey in SSHattacks.keys():
                SSHattacks [AttackKey] = SSHattacks [AttackKey] + 1
            else:
                SSHattacks [AttackKey] = 1

    if psummary:
        print ("\nSSH Success Count :", SSHsuccess)
        print ("SSH Failure Count :", SSHfail)

    for key, value in SSHattacks.items():
        print("\n", key, ' : ', value)



def SaveSessionsToFile(sessions):
#Saves individual sessions to PCAP files

    if len(sessions) != 0:
        Keys = list(sessions.keys())
        Values = list(sessions.values())

        print("\nSaving Sessions to file...")

        for Kindex,Vindex in zip(Keys, Values):
            Key = Kindex.replace("[","").replace("]","").replace("'","").replace(", ","-") + ".pcap"

            wrpcap(Key,Vindex)



def main():

    Capture=sniff(offline="sshbrute.pcap")
    #Capture=sniff(offline="ftpbrute.pcap")

    PrintSummaryToConsole = True #Print out the summary of analysis
    PrintLoginsToConsole = False #Print out the logins found
    SaveToFile = True #Save packetlists containing failures
    TrackAttackSources = True #Save failure sources and destination and print individual summary
    SSHThreshold = 7500


    FullDuplexSessions = Capture.sessions(Full_Duplex_Sessions) #Create a dict of full duplex session streams

    TCPsessions = TCPSessionFilter (FullDuplexSessions)

    #FTPSessions = FTPSessionFilter (TCPsessions) #Only returns TCP/FTP sessions

    #if len(FTPSessions) != 0: #Do any FTP sessions exist?
    #    FTPBruteAnalysis (FTPSessions,PrintLoginsToConsole,PrintSummaryToConsole,SaveToFile,TrackAttackSources)



    SSHSessions = SSHSessionFilter (TCPsessions)

    if len(SSHSessions) != 0: #Do any SSH sessions exist?
        SSHBruteAnalysis (SSHSessions, SSHThreshold, PrintSummaryToConsole, TrackAttackSources)



if __name__ == "__main__":
    main()


