
import logging
import sys
from netaddr import *

class PcapEngine():

    def __init__(self, pcap_file_name, pcap_parser_engine="scapy"):
        self.packetDB = {}
        self.destination_address = {}
        if pcap_parser_engine == "scapy":
            try:
                from scapy.all import rdpcap
            except:
                logging.error("Cannot import selected pcap engine: Scapy!")
                sys.exit()
            # Scapy sessions and other types use more O(N) iterations so just
            # use rdpcap
            self.packets = rdpcap(pcap_file_name)
        elif pcap_parser_engine == "pyshark":
            try:
                import pyshark
            except:
                logging.error("Cannot import selected pcap engine: PyShark!")
                sys.exit()
            self.packets = pyshark.FileCapture(pcap_file_name)
        # Add other pcap engine modules to generate packetDB
        self.analyse_packet_data()

    def analyse_packet_data(self):
        """
        PcapXray runs only one O(N) packets once to memoize
        # - Parse the packets to create a usable DB
        # - Store the db as json on a file in cache folder (to not repeat read)
        # - All the protocol parsing should be included here
        """
        for packet in self.packets: # O(N) packet iteration 
            if "IP" in packet:
                # Handle IP packets that originated from LAN (Internal Network)
                if IPAddress(packet["IP"].src).is_private():
                    key = packet["IP"].src + " => " + packet["IP"].dst + ":" + packet["TCP"].dport if "TCP" in packet else packet["UDP"].dport
                    source_private_ip = key
                    if source_private_ip not in self.packetDB:
                        self.packetDB[source_private_ip] = {}
                        # TCP
                        if "TCP" in packet and "TCP" not in self.packetDB[source_private_ip]:
                            self.packetDB[source_private_ip]["TCP"] = {}
                        # UDP 
                        if "UDP" in packet and "UDP" not in self.packetDB[source_private_ip]:
                            self.packetDB[source_private_ip]["UDP"] = {}
                        # Ethernet Layer ( Mac address )
                        if "Ether" in packet and "Ethernet" not in self.packetDB[source_private_ip]:
                            self.packetDB[source_private_ip]["Ethernet"] = packet["Ethernet"].src
                        # HTTP Packets
                        if "TCP" in packet and packet["TCP"].dport == 80:
                            if "HTTP" not in self.packetDB[source_private_ip]["TCP"]:
                                self.packetDB[source_private_ip]["TCP"]["HTTP"] = {}
                            if "Server" not in self.packetDB[source_private_ip]["TCP"]["HTTP"]:
                                self.packetDB[source_private_ip]["TCP"]["HTTP"]["Server"] = []
                            if "Payload" not in self.packetDB[source_private_ip]["TCP"]["HTTP"]:
                                self.packetDB[source_private_ip]["TCP"]["HTTP"]["Payload"] = []
                            # Destination Address
                            self.packetDB[source_private_ip]["TCP"]["HTTP"]["Server"].append(packet["IP"].dst)
                            # Server Details
                            self.packetDB[source_private_ip]["TCP"]["HTTP"]["Server"] = list(set(self.packetDB[packet["IP"].src]["TCP"]["HTTP"]["Server"]))
                            # Payload recording
                            self.packetDB[source_private_ip]["TCP"]["HTTP"]["Payload"].append(packet)
                        # SSL Packets
                        # - Get handshake details?
                        if "TCP" in packet and packet["TCP"].dport == 443:
                            if "HTTPS" not in self.packetDB[source_private_ip["TCP"]:
                                self.packetDB[source_private_ip]["TCP"]["HTTPS"] = []
                            if packet["IP"].dst not in self.packetDB[source_private_ip]["TCP"]["HTTPS"]:
                                self.packetDB[source_private_ip]["TCP"]["HTTPS"].append(packet["IP"].dst)
                                self.packetDB[source_private_ip]["TCP"]["HTTPS"] = list(set(self.packetDB[source_private_ip]["TCP"]["HTTPS"]))
                            if "TCP" in packet:
                                if "PortsConnected" not in self.packetDB[source_private_ip]["TCP"]:
                                    self.packetDB[source_private_ip]["TCP"]["PortsConnected"] = []
                                destinaton_port = packet["TCP"].dport
                                destination_ip = packet["IP"].dst
                                if (destination_ip,destination_port) not in self.packetDB[source_private_ip]["TCP"]["PortsConnected"]:
                                    self.packetDB[source_private_ip]["TCP"]["PortsConnected"].append((destination_ip,destination_port))
                            if packet.haslayer(UDP):
                                if "PortsConnected" not in self.packetDB[source_private_ip]["UDP"]:
                                    self.packetDB[source_private_ip]["UDP"]["PortsConnected"] = []
                                destination_port = packet.getlayer(UDP).dport
                                destination_ip = packet.getlayer(IP).dst
                                if (destination_ip,destination_port) not in self.packetDB[source_private_ip]["UDP"]["PortsConnected"]:
                                    self.packetDB[source_private_ip]["UDP"]["PortsConnected"].append((destination_ip,destination_port))

                            #HTTPS
                            #Tor
                            #Malicious
                            # HTTP Payload Decrypt
                        if IPAddress(packet.getlayer(IP).dst).is_private():
                            if packet.getlayer(IP).dst not in self.packetDB:
                                self.packetDB[packet.getlayer(IP).dst] = {}
                            if packet.haslayer(TCP) and "TCP" not in self.packetDB[packet.getlayer(IP).dst]:
                                self.packetDB[packet.getlayer(IP).dst]["TCP"] = {}
                            if packet.haslayer(UDP) and "UDP" not in self.packetDB[packet.getlayer(IP).dst]:
                                self.packetDB[packet.getlayer(IP).dst]["UDP"] = {}
                            if packet.haslayer(Ether) and "Ethernet" not in self.packetDB[packet.getlayer(IP).dst]:
                                self.packetDB[packet.getlayer(IP).dst]["Ethernet"] = packet.getlayer(Ether).dst
                            if packet.haslayer(TCP) and packet.getlayer(TCP).sport == 80:
                                if "HTTP" not in self.packetDB[packet.getlayer(IP).dst]["TCP"]:
                                    self.packetDB[packet.getlayer(IP).dst]["TCP"]["HTTP"] = {}
                                if "Server" not in self.packetDB[packet.getlayer(IP).dst]["TCP"]["HTTP"]:
                                    self.packetDB[packet.getlayer(IP).dst]["TCP"]["HTTP"]["Server"] = []
                                if "Payload" not in self.packetDB[packet.getlayer(IP).dst]["TCP"]["HTTP"]:
                                    self.packetDB[packet.getlayer(IP).dst]["TCP"]["HTTP"]["Payload"] = []
                                self.packetDB[packet.getlayer(IP).dst]["TCP"]["HTTP"]["Server"].append(packet.getlayer(IP).src)
                                self.packetDB[packet.getlayer(IP).dst]["TCP"]["HTTP"]["Server"] = list(set(self.packetDB[packet.getlayer(IP).dst]["TCP"]["HTTP"]["Server"]))
                                self.packetDB[packet.getlayer(IP).dst]["TCP"]["HTTP"]["Payload"].append(packet)
                            if packet.haslayer(TCP) and packet.getlayer(TCP).sport == 443:
                                if "HTTPS" not in self.packetDB[packet.getlayer(IP).dst]["TCP"]:
                                    self.packetDB[packet.getlayer(IP).dst]["TCP"]["HTTPS"] = []
                                if packet.getlayer(IP).src not in self.packetDB[packet.getlayer(IP).dst]["TCP"]["HTTPS"]:
                                    self.packetDB[packet.getlayer(IP).dst]["TCP"]["HTTPS"].append(packet.getlayer(IP).src)
                                    self.packetDB[packet.getlayer(IP).dst]["TCP"]["HTTPS"] = list(set(self.packetDB[packet.getlayer(IP).dst]["TCP"]["HTTPS"]))
                            if packet.haslayer(TCP):
                                if "PortsConnected" not in self.packetDB[packet.getlayer(IP).dst]["TCP"]:
                                    self.packetDB[packet.getlayer(IP).dst]["TCP"]["PortsConnected"] = []
                                port = packet.getlayer(TCP).sport
                                ip = packet.getlayer(IP).src
                                if (ip,port) not in self.packetDB[packet.getlayer(IP).dst]["TCP"]["PortsConnected"]:
                                    self.packetDB[packet.getlayer(IP).dst]["TCP"]["PortsConnected"].append((ip,port))
                            if packet.haslayer(UDP):
                                if "PortsConnected" not in self.packetDB[packet.getlayer(IP).dst]["UDP"]:
                                        self.packetDB[packet.getlayer(IP).dst]["UDP"]["PortsConnected"] = []
                                port = packet.getlayer(UDP).sport
                                ip = packet.getlayer(IP).src
                                if (ip, port) not in self.packetDB[packet.getlayer(IP).dst]["UDP"]["PortsConnected"]:
                                    self.packetDB[packet.getlayer(IP).dst]["UDP"]["PortsConnected"].append((ip, port))

        
# Module Driver
def main():
    pcapfile = PcapEngine('examples/test.pcap', "pyshark")
    #print pcapfile.packetDB
    #print(self.packetDB["TCP 192.168.0.26:64707 > 172.217.12.174:443"].summary())
    #print(self.packetDB["TCP 172.217.12.174:443 > 192.168.0.26:64707"].summary())
    #self.packetDB.conversations(type="jpg", target="> test.jpg")

main()










import logging
import sys
from netaddr import *

class PcapEngine():

    def __init__(self, pcap_file_name, pcap_parser_engine="scapy"):
        """
        PcapEngine: To support different pcap parser backend engine to operate reading pcap
        Current Support:
        * Scapy
        * Pyshark
        - The init function imports libraries based on the parser engine selected
        Return:
        * packetDB ==> Full Duplex Packet Streams
          - Used while finally plotting streams as graph
          - dump packets during report generation
        * lan_hosts ==> Private IP (LAN) list
          - device details
        * destination_hosts ==> Destination Hosts
          - communication details
          - tor identification
          - malicious identification
        """
        self.packetDB = {}
        self.lan_hosts = {}
        self.destination_hosts = {}
        self.engine = pcap_parser_engine
        if pcap_parser_engine == "scapy":
            try:
                from scapy.all import rdpcap
            except:
                logging.error("Cannot import selected pcap engine: Scapy!")
                sys.exit()
            # Scapy sessions and other types use more O(N) iterations so just
            # - use rdpcap + our own iteration (create full duplex streams)
            self.packets = rdpcap(pcap_file_name)
        elif pcap_parser_engine == "pyshark":
            try:
                import pyshark
            except:
                logging.error("Cannot import selected pcap engine: PyShark!")
                sys.exit()
            self.packets = pyshark.FileCapture(pcap_file_name)
        
        # Pyshark has different TCP port access change + Ethernet layer access 
        self.analyse_packet_data_with_pyshark()

        # Add other pcap engine modules to generate packetDB
        

    def analyse_packet_data_with_scapy(self):
        """
        PcapXray runs only one O(N) packets once to memoize
        # - Parse the packets to create a usable DB
        # - Store the db as json on a file in cache folder (to not repeat read)
        # - All the protocol parsing should be included here
        """
        for packet in self.packets: # O(N) packet iteration 
            if "IP" in packet:
                # Handle IP packets that originated from LAN (Internal Network)
                source_private_ip = None
                if "TCP" in packet or "UDP" in packet:
                    if IPAddress(packet["IP"].src).is_private():
                        key = packet["IP"].src + "/" + packet["IP"].dst + "/" + str(packet["TCP"].dport) if "TCP" in packet else str(packet["UDP"].dport)
                        source_private_ip = key
                    elif IPAddress(packet["IP"].dst).is_private():
                        key = packet["IP"].dst + "/" + packet["IP"].src + "/" + str(packet["TCP"].sport) if "TCP" in packet else str(packet["UDP"].sport)
                        source_private_ip = key
                elif "ICMP" in packet:
                    key = packet["IP"].src + "/" + packet["IP"].dst + "/" + "ICMP"
                    source_private_ip = key
                
                # IntraNetwork vs InterNetwork Hosts list
                if private_source:
                    self.lan_hosts[packet["IP"].src] = ""
                else:
                    self.destination_hosts[packet["IP"].src] = {}
                if private_destination:
                   self.lan_hosts[packet["IP"].dst] = ""
                else:
                    self.destination_hosts[packet["IP"].dst] = {}
                
            if source_private_ip:
                if source_private_ip not in self.packetDB:
                    self.packetDB[source_private_ip] = {}
                    # Ethernet Layer ( Mac address )
                    if "Ether" in packet and "Ethernet" not in self.packetDB[source_private_ip]:
                        self.packetDB[source_private_ip]["Ethernet"] = {}
                    self.packetDB[source_private_ip]["Ethernet"]["src"] = packet["Ethernet"].src
                    self.packetDB[source_private_ip]["Ethernet"]["dst"] = packet["Ethernet"].dst
                    # HTTP Packets
                    if "Payloads" not in packet:
                        # Payload recording
                        self.packetDB[source_private_ip]["Payload"] = []
                    self.packetDB[source_private_ip]["Payload"].append(packet["Raw"].load if "Raw" in packet else None) 
        
    def analyse_packet_data_with_pyshark(self):
        """
        PcapXray runs only one O(N) packets once to memoize
        # - Parse the packets to create a usable DB
        # - Store the db as json on a file in cache folder (to not repeat read)
        # - All the protocol parsing should be included here
        """
        for packet in self.packets: # O(N) packet iteration 
            if "IP" in packet:
                # Handle IP packets that originated from LAN (Internal Network)
                source_private_ip = None
                private_source = IPAddress(packet["IP"].src).is_private()
                private_destination = IPAddress(packet["IP"].src).is_private()

                if "TCP" in packet or "UDP" in packet:
                     # Sort out indifferences in pcap engine
                    if self.engine == "pyshark":
                        tcp_src = str(packet["TCP"].srcport if "TCP" in packet else packet["UDP"].srcport)
                        tcp_dst = str(packet["TCP"].dstport if "TCP" in packet else packet["UDP"].dstport)
                    else:
                        tcp_src = str(packet["TCP"].sport if "TCP" in packet else packet["UDP"].sport)
                        tcp_dst = str(packet["TCP"].dport if "TCP" in packet else packet["UDP"].dport)

                    if private_source:
                        key = packet["IP"].src + "/" + packet["IP"].dst + "/" + tcp_dst
                        source_private_ip = key
                    elif private_destination:
                        key = packet["IP"].dst + "/" + packet["IP"].src + "/" + tcp_src
                        source_private_ip = key

                elif "ICMP" in packet:
                    key = packet["IP"].src + "/" + packet["IP"].dst + "/" + "ICMP"
                    source_private_ip = key
                
                # IntraNetwork vs InterNetwork Hosts list
                if private_source and not private_destination:
                    self.lan_hosts[packet["IP"].src] = ""
                else:
                    self.destination_hosts[packet["IP"].src] = {}
                if private_destination and not private_source:
                   self.lan_hosts[packet["IP"].dst] = ""
                else:
                    self.destination_hosts[packet["IP"].dst] = {}
                
            if source_private_ip:
                if source_private_ip not in self.packetDB:
                    self.packetDB[source_private_ip] = {}
                    # Ethernet Layer ( Mac address )
                    if "Ethernet" not in self.packetDB[source_private_ip]:
                        self.packetDB[source_private_ip]["Ethernet"] = {}
                    # HTTP Packets
                    if "Payloads" not in self.packetDB:
                        # Payload recording
                        self.packetDB[source_private_ip]["Payload"] = []
                if self.engine == "pyshark":
                    self.packetDB[source_private_ip]["Ethernet"]["src"] = packet["ETH"].src
                    self.packetDB[source_private_ip]["Ethernet"]["dst"] = packet["ETH"].dst
                    if "data" in packet: 
                        self.packetDB[source_private_ip]["Payload"].append(packet.data)
                else:
                    self.packetDB[source_private_ip]["Ethernet"]["src"] = packet["Ethernet"].src
                    self.packetDB[source_private_ip]["Ethernet"]["dst"] = packet["Ethernet"].dst
                    if "Raw" in packet:
                        self.packetDB[source_private_ip]["Payload"].append(packet["Raw"].load)

# Module Driver
def main():
    pcapfile = PcapEngine('examples/test.pcap', "pyshark")
    print(pcapfile.packetDB)
    print(pcapfile.lan_hosts)
    print(pcapfile.destination_hosts)
    #print(self.packetDB["TCP 192.168.0.26:64707 > 172.217.12.174:443"].summary())
    #print(self.packetDB["TCP 172.217.12.174:443 > 192.168.0.26:64707"].summary())
    #self.packetDB.conversations(type="jpg", target="> test.jpg")

main()

# Sort payload by time...
# SSL Packets
# - Get handshake details?
#HTTPS
#Tor
#Malicious
# HTTP Payload Decrypt
