import pyaiengine

if __name__ == '__main__':

    # Load an instance of a Lan Network Stack
    st = pyaiengine.StackLan()

    # Allocate the UDP flows in order to keep memory
    # under control and avoid allocations during the execution
    st.udp_flows = 163840

    with pyaiengine.PacketDispatcher("./pcapfile/somefile.pcap") as pd:
        pd.stack = st
        pd.run()

    # Get the UDP flows processed
    flows = st.udp_flow_manager
    print("Total flows:%d", len(flowlist))

    domains = dict()
    for flow in flows:
        if(flow.bytes  > 0):
            if (flow.dns_info):
                name = flow.dns_info.domain_name
                if(not domains.has_key(name)):
                    domains[name] = 1
                else:
                    domains[name] += 1
    print(domains)
    sys.exit(0)

import pyaiengine

def callback_drop_packets(flow_name):
    """ Send a command to the Iptables in other to drop the packets """
    source_ip = str(flow_name).split(":")[0]
    os.system("iptables -A INPUT -i eth0 -s %s -j DROP" % flow.src_ip)

def loadSignaturesForTcp():
    """ Load the signatures from source, Snort, Suricata, etc. """

     sm = pyaiengine.RegexManager()

     sig = pyaiengine.Regex("Shellcode Generic Exploit","\x90\x90\x90\x90\x90\x90\x90\x90\x90")

     """ Sets a specific callback to the signature created """
     sig.calllback = callback_drop_packets
     sm.add_regex(sig)

     return sm

if __name__ == '__main__':

     # Load an instance of a Network Stack
     st = pyaiengine.StackLan()

     # Load Signatures/Rules in order to detect the traffic
     s_tcp = loadSignaturesForTcp()
     st.tcp_regex_manager = s_tcp

     st.tcp_flows = 327680
     st.udp_flows = 163840

     with pyaiengine.PacketDispatcher("eth0") as pd:
        pd.stack = st 
        pdis.run()

import pyaiengine

def callback_domain(flow):

    ip = str(flow).split(":")[0]
    domain = flow.dns_info
    if(domain):
        print("Malware on ip %s domain %s" % (ip,domain.domain_name)) 

def loadBadDomains():

    dm = pyaiengine.DomainNameManager()

    # List from http://www.malwaredomainlist.com/hostslist/hosts.txt
    # http://www.abuse.ch/zeustracker/blocklist.php?download=domainblocklist for ZeusDomains
    # Parse the file and add the domains.
    f = open("hosts.txt","r")

    lines = f.readlines()
    i = 0
    for line in lines:
        if( line[0] != "#"):
            domain = line.replace("\r\n","").split(" ")
            if(len(domain)> 2):
                name = "Bad domain %d" % i
                i = i +1
                dom = pyaiengine.DomainName(name,domain[2])
                dom.callback = callback_domain

                dm.add_domain_name(dom)
    f.close()
    return dm

if __name__ == '__main__':

     # Load an instance of a Network Stack on Mobile network
     st = pyaiengine.StackMobile()

     doms = loadBadDomains()
     st.set_domain_name_manager(doms,"DNSProtocol")

     # 1.638.400 udp flows will be enough?
     st.udp_flows = 1638400

     with pyaiengine.PacketDispatcher("eth0") as pd:
         pd.stack = st
         pd.run()

import pyaiengine

r_mng = None

def callback_troyan_activity(flow):
    ip = str(flow).split(":")[0]

    print("Detected OSX_DocksterTrojan on ip:%s" % flow.src_ip)

def callback_domain(flow):

    ip = str(flow).split(":")[0]

    print("Suspicious Domain (%s) from %s" % (flow.dns_info.domain_name,ip))
    print("Add specific regex for OSX_DocksterTrojan")

    reg = pyaiengine.Regex("OSX_DocksterTrojan regex activity",
        "^\xff\xff\xff\xff\xc2\x1f\x96\x9b\x5f\x03\xd3\x3d\x43\xe0\x4f\x8f")

    reg.callback = callback_troyan_activity
    r_mng.add_regex(reg)

    st.tcp_regex_manager = r_mng


if __name__ == '__main__':

    # Load an instance of a Network Stack on Mobile network
    st = pyaiengine.StackMobile()

    dm = pyaiengine.DomainNameManager()

    dom = pyaiengine.DomainName("OSX_DocksterTrojan suspicious domain",
        "itsec.eicp.net")
    dom.callback = callback_domain
    dm.add_domain_name(dom)

    st.set_domain_name_manager(dm,"DNSProtocol")

    st.tcp_flows = 327680
    st.udp_flows = 163840


    with pyaiengine.PacketDispatcher("eth1") as pd:
        pd.stack = st    
        pd.run()

import pyaiengine
import redis

class redisAdaptor(pyaiengine.DatabaseAdaptor):
    """ This class inheritance of DatabaseAdaptor that contains 
    the following methods:
        - insert, called on the first insertion of the network flow
        - update, called depending on the sample selected.
        - delete, called when the flow is destroy.
    """
    def __init__(self):
        self.__r = None 
        self.__total_inserts = 0
        self.__total_updates = 0
        self.__total_removes = 0

    def connect(self,connection_str):
       self.__r = redis.Redis(connection_str)   

    def update(self,key,data):
        self.__r.hset("udpflows",key,data)
        self.__total_updates = self.__total_updates + 1 

    def insert(self,key):
        self.__r.hset("udpflows",key,"{}")
        self.__total_inserts = self.__total_inserts + 1

    def delete(self,key):
        self.__r.hdelete("udpflows",key)
        self.__total_removes = self.__total_removes + 1

    def show(self):
        print("Total inserts %d, total updates %d, total removes %d" % (self.__total_inserts,self.__total_updates,self.__total_removes))

if __name__ == '__main__':

    # Load an instance of a Network Stack on Mobile network
    st = pyaiengine.StackLan()

    st.tcp_flows = 327680
    st.udp_flows = 163840

    """
    Create a redisAdaptor object. 
    This is just and example you can create your own adaptor for
    any database.
    """
    db = redisAdaptor()
    # connect to the redis database 
    db.connect("localhost")

    """ 
    Set the database adaptor just for UDP traffic
        and with a packet sampling of 512 packets, so every 512 packets
        the method "update" will be called.
        Fix this value depending on your software/hardware requirments.
    """
    st.set_udp_database_adaptor(db,512)

    with pyaiengine.PacketDispatcher("eth0") as pd:
        pd.stack = st
        pd.run()

    db.show()
    sys.exit(0)

import pyaiengine

def callback_heartbeat(flow):

    payload = flow.payload
    if(len(payload) > 7): 
        """ Heartbeat minimum header """
        if(int(payload[7])>1):
            print("SSL Heartbeat leak on %s" % str(flow))
            print(payload)

if __name__ == '__main__':

    # Load an instance of a Network Stack
    st = pyaiengine.StackLan()

    sm = pyaiengine.RegexManager()

    """ 
    Heartbeat regex expression
    18 -> Content Type: Heartbeat
        0301, 0302 -> Version: TLS
        xxxx -> Length
        01 - Heartbeat
        xx - heartbeat payload length
    """ 
    sig = pyaiengine.Regex("SSL Heartbeat","^\x18\x03(\x01|\x02|\x03)\x00\x03\x01")
    sig.callback = callback_heartbeat
    sm.add_regex(sig)

    st.tcp_regex_manager = sm

    st.tcp_flows = 327680
    st.udp_flows = 163840

    st.enable_nids_engine = True

    with pyaiengine.PacketDispatcher("eth0") as pd:
        pd.stack = st
        pd.run()

    sys.exit(0)

import pyaiengine
from netfilterqueue import NetfilterQueue

""" Need a fake ethernet header """
ethernet_header = "\xbe\xef\x00\x00\x00\x01\xbe\xef\x00\x00\x00\x02\x08\x00"

# Create a instace of a PacketDispatcher
pdis = pyaiengine.PacketDispatcher()

def netfilter_callback(packet):

    payload = ethernet_header + packet.get_payload()
    length = packet.get_payload_len() + 14

    """ Use the forwardPacket method from the PacketDispatcher object
    in order to forward the packets from netfilter """
    pdis.forward_packet(payload,length)
    packet.accept()

if __name__ == '__main__':

    # Load an instance of a Network Stack
    st = pyaiengine.StackLan()

    # Plug the stack on the PacketDispatcher
    pdis.stack = st

    st.tcp_flows = 327680
    st.udp_flows = 163840

    """ Create a NetfilterQueue object """
    nfqueue = NetfilterQueue()

    """ Sets the callback for netfilter """
    nfqueue.bind(1, netfilter_callback)
    try:
        nfqueue.run()
    except KeyboardInterrupt:
        pass

    sys.exit(0)