#!/usr/bin/python3

from setproctitle import setproctitle
from shell_cmd import sh
from nfstream import NFStreamer
import time
import sys, signal
import json
import logging


logging.basicConfig(filename="/tmp/dpi.log",
                    filemode='a',
                    format='%(asctime)s,%(msecs)d %(name)s %(levelname)s %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S',
                    level=logging.WARNING)

log=logging.getLogger("NexDPI")

deftimeout="3600"

setproctitle("dpi")

online_streamer = NFStreamer(source="eth1", promiscuous_mode=False, splt_analysis=20, statistical_analysis=False)
 
templconf = """
{
  "Log": "ERROR",

  "Cats":{

     "Network":{
        "nostart": [],
        "noapps": [],
        "knownapps": ["DHCP", "DHCPV6", "DNScrypt", "DoH_DoT", "DNS", "Ookla", "ICMP", "ICMPV6", "IGMP", "LLMNR", "MDNS", "DoH_DoT", "Ookla", "WSD"],
        "knowstarts": [],
        "ipset": "system_triplet",
        "timeout":"3600"
     },

     "Game":{
        "nostart": [],
        "noapps": [],
        "knownapps": ["Steam", "Xbox", "Playstation"],
        "knowstarts": ["TLS", "HTTP"],
        "ipset": "streaming_triplet",
        "timeout":"3600"
     },

     "Template-GroupName":{
        "nostart": ["DNS", "ICMP"],
        "noapps": [],
        "knownapps": [],
        "knowstarts": [],
        "ipset": "name_ipset",
        "timeout":"3600"
     }
  },

  "Apps":{
     "TikTok":{
        "ipset": "kids_triplet",
        "nostart": [],
        "timeout": "3600",
        "knowstarts":"TLS"
     }
  },

  "Ignore": [
     ["DHCPv6", "Network"],
     ["DHCP", "Network"]
  ]

}
"""

try:
   fconf = open("/etc/nexdpi/dpirules.json", "r")
   R=json.loads(fconf.read())
   fconf.close()
except:
   print(" Cant start without a rules file.\n\n")
   print(" Please use the following example to create a JSON config file as /etc/nexdpi/dpirules.json\n\n")
   print(templconf)
   sys.exit(0)


Cats = R['Cats']
Apps = R['Apps']
Ignore = list(R['Ignore'])
Config = {
      'throttle': 0
      }

if 'Config' in R.keys():
   Config = R['Config']

if 'Log' in R.keys():
   if R['Log'] == 'DEBUG':
      log.setLevel(logging.DEBUG)
   elif R['Log'] == "INFO":
      log.setLevel(logging.INFO)
   elif R['Log'] == "WARNING":
      log.setLevel(logging.WARNING)
   elif R['Log'] == "ERROR":
      log.setLevel(logging.ERROR)


def reloadconf(signum, frame):
   global Cats
   global Apps
   global Ignore
   global Config

   try:
      fconf = open("/etc/nexdpi/dpirules.json", "r")
      R=json.loads(fconf.read())
      fconf.close()
      Cats = R['Cats']
      Apps = R['Apps']
      Ignore = list(R['Ignore'])
      if 'Log' in R.keys():
         if R['Log'] == 'DEBUG':
            log.setLevel(logging.DEBUG)
         elif R['Log'] == "INFO":
            log.setLevel(logging.INFO)
         elif R['Log'] == "WARNING":
            log.setLevel(logging.WARNING)
         elif R['Log'] == "ERROR":
            log.setLevel(logging.ERROR)
      if 'Config' in R.keys():
         Config = R['Config']

      log.warning("Rules file reloaded")
   except:
      log.error("Error loading rules file.")



signal.signal(signal.SIGHUP, reloadconf)

UnknownMatch=[]


"""
NFlow(id=5,
      expiration_id=0, 
      src_ip=192.168.42.140,
      src_mac=ee:9a:15:8d:4c:2e,
      src_oui=ee:9a:15,
      src_port=37860,  
      dst_ip=102.132.99.1,
      dst_mac=40:62:31:05:c8:56,
      dst_oui=40:62:31,
      dst_port=443,
      protocol=6,
      ip_version=4,
      vlan_id=0,
      bidirectional_first_seen_ms=1619278659528,
      bidirectional_last_seen_ms=1619278660194,
      bidirectional_duration_ms=666,
      bidirectional_packets=25,
      bidirectional_bytes=10856,
      src2dst_first_seen_ms=1619278659528,
      src2dst_last_seen_ms=1619278660194,
      src2dst_duration_ms=666,
      src2dst_packets=14,
      src2dst_bytes=2930,
      dst2src_first_seen_ms=1619278659532,
      dst2src_last_seen_ms=1619278660193,
      dst2src_duration_ms=661,
      dst2src_packets=11,
      dst2src_bytes=7926,
      splt_direction=[0, 1, 0, 0, 1, 1, 0, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0, 0, 1, 1],
      splt_ps=[74, 74, 66, 583, 66, 3309, 66, 66, 66, 130, 66, 243, 721, 66, 66, 1300, 66, 420, 66, 1297],
      splt_piat_ms=[0, 4, 1, 11, 2, 1, 1, 0, 0, 22, 2, 0, 1, 2, 37, 126, 1, 1, 3, 167],
      application_name=TLS.Facebook,
      application_category_name=SocialNetwork,
      application_is_guessed=0,
      requested_server_name=graph.facebook.com,
      client_fingerprint=9b02ebd3a43b62d825e1ac605b621dc8,
      server_fingerprint=,
      user_agent=,
      content_type=)   
"""


class NexDPI():

   fullname=False

   def main(self):
      log.warning("NexDPI started")
      for flow in online_streamer:

         managed=False
         self.fullname=flow.application_name+" "+flow.application_category_name
         triplet=str(flow.dst_ip)+","+str(flow.dst_port)+","+str(flow.src_ip)
         cname = flow.application_category_name
         aname = flow.application_name
         sername = aname.split(".")[-1:][0]
         startname = False
         if aname != sername:
            startname =  aname.split(".")[:-1][0]
         ipv=flow.ip_version

         log.debug("RECEIVED: "+cname+" "+aname+" "+sername+"\n\n"+str(flow))

         if [aname, cname] in list(Ignore):
            log.debug("IGNORED: "+self.fullname)
            continue

         if sername in list(Apps.keys()):
            appd=Apps[sername]
            if ipv==6:
               ipset_list = appd['ipset']+"6"
            else:
               ipset_list = appd['ipset']
            if not aname.startswith(tuple(appd['nostart'])):
               managed=ipset_list+" Apps"
               sh("ipset test "+ipset_list+" "+triplet+" >/dev/null 2>&1 || ipset add "+ipset_list+" "+triplet+" timeout "+appd['timeout']+" > /dev/null 2>&1")
               log.info("ADD: "+ipset_list+" "+triplet+" "+self.fullname)
               if aname.startswith(tuple(appd['knowstarts'])):
                  continue

         if cname in list(Cats.keys()):
            if ipv==6:
               ipset_list = Cats[cname]['ipset']+"6"
            else:
               ipset_list = Cats[cname]['ipset']
            if not aname.startswith(tuple(Cats[cname]['nostart'])) and not sername in list(Cats[cname]['noapps']):
               managed=ipset_list+" Cats"
               sh("ipset add "+ipset_list+" "+triplet+" timeout "+Cats[cname]['timeout']+" --exist > /dev/null 2>&1")
               log.info("ADD: "+ipset_list+" "+triplet+" "+self.fullname)
               if sername in list(Cats[cname]['knownapps']):
                  if not startname:
                     continue
                  elif startname in list(Cats[cname]['knowstarts']):
                     continue

         if self.fullname and not self.fullname in UnknownMatch:
            log.warning("UNKNOWN("+str(managed)+"): "+self.fullname)
            f = open("/tmp/dpi.unknown", "a")
            f2 = open("/tmp/dpi.ignore.unknown", "a")
            f.write(aname+" "+cname+"\n")
            f2.write("    [\""+aname+"\",\""+cname+"\"],\n")
            f.close()
            f2.close()
            if managed:
               log.warning("MANAGED_UNKNOWN: "+self.fullname)
               f = open("/tmp/dpi.managed.unknown", "a")
               f.write(aname+" "+cname+" "+managed+"\n")
               f.close()

            UnknownMatch.append(self.fullname)

         if Config['throttle'] > 0:
            time.sleep(Config['throttle'])



if __name__ == "__main__":
   ndpi=NexDPI()
   ndpi.main()
   
