from scapy.all import *
import datetime
import os
import binascii
import yaml
import sqlite3
from datetime import datetime
import protocol_actions
import threading
import requests
lock = threading.Lock()
from helpers.Packet_Analyzer import Packet_Analyzer

os.chdir('.')
mydb=sqlite3.connect("assets.db",check_same_thread=False)

try:
	yamlfile=open("config.yaml")
	data = yaml.load(yamlfile, Loader=yaml.FullLoader)
	server_ip=data["configuration"]["unifytwin_server_ip_address"]
	packet_storage=data["configuration"]["packet_storage"]
	plant=data["configuration"]["edge_device_location"]
	balacklist_ips=data["configuration"]["blacklist_ip"]
	blacklist_dns=data["configuration"]["blacklist_dns"]
	configured_threads=data["configuration"]["threads"]
	interfaces=data["configuration"]["interfaces"]
	null_loopback=data["configuration"]["null_loopback"]
except Exception as e:
	server_ip=""
	packet_storage="local"
	plant="Unknown Location"
	balacklist_ips=['46.4.105.116','172.67.214.157','3.6.115.64','104.21.53.154']
	blacklist_dns=['webhook.site.','hpd.gasmi.net.','dd01-14-98-12-178.in.ngrok.io.']
	configured_threads=10
	interfaces='eth0'
	null_loopback=False


from sys import platform
if platform == "linux" or platform == "linux2":
    clearing='clear'
elif platform == "darwin":
    clearing='clear'
elif platform == "win32":
    clearing='cls'


def convert_text(pkt):
	query=mydb.cursor()
	lock.acquire()
	lock.release()
	inhex=binascii.hexlify(bytes(pkt))
	if(null_loopback):
		inhex=b'0000000000000000000000000800'+inhex[8:] if inhex[:2]==b'02' else b'00000000000000000000000086DD'+inhex[8:]
	try:
		pa = Packet_Analyzer()
		protocols=pa.identify_protocol(inhex)
		ip=pa.get_ip(inhex)
		if(ip):
			mac=pa.get_mac(inhex)
			check_exist=query.execute("SELECT ip_address FROM inventory WHERE ip_address='"+ip+"'").fetchone()
			if(not check_exist):
				dev_type,vendor,firmware,model=protocol_actions.analyse_protocol(protocols,pkt)
				insinventory='INSERT INTO inventory\
				(ip_address,name,last_activity,type,protocols,mac_address,vendor,firmware_version,model,operating_system,plant)\
				VALUES\
				("'+ip+'","'+ip+'","'+datetime.now().strftime("%d-%m-%y %H:%M:%S")+'","'+dev_type+'","'+str(protocols)+'","'+mac+'","'+vendor+'","'+firmware+'","'+model+'","Unknown","'+plant+'")'
				query.execute(insinventory)
				query.close()
				mydb.commit()
			else:
				prev_protocols=query.execute("SELECT protocols FROM inventory WHERE ip_address='"+ip+"'").fetchone()[0]
				extra_proto=set(protocols.split(':'))-set(prev_protocols.split(':'))
				if(extra_proto):
					dev_type,vendor,firmware,model=protocol_actions.update_protocol(protocols,pkt)
					prev_protocols=prev_protocols+':'+str(extra_proto).replace(", ",":").replace("{","").replace("}","").replace("'","")
					if(dev_type!="Unknown"):
						updateproto='UPDATE inventory SET protocols="'+prev_protocols+'",last_activity="'+datetime.now().strftime("%d-%m-%y %H:%M:%S")+'",type="'+dev_type+'",vendor="'+vendor+'",firmware_version="'+firmware+'",model="'+model+'" WHERE ip_address="'+ip+'"'
						query.execute(updateproto)
						query.close()
						mydb.commit()
					else:
						updateproto='UPDATE inventory SET protocols="'+prev_protocols+'",last_activity="'+datetime.now().strftime("%d-%m-%y %H:%M:%S")+'" WHERE ip_address="'+ip+'"'
						query.execute(updateproto)
						query.close()
						mydb.commit()
				else:
					update_last_activity='UPDATE inventory SET last_activity="'+datetime.now().strftime("%d-%m-%y %H:%M:%S")+'" WHERE ip_address="'+ip+'"'
					query.execute(update_last_activity)
					query.close()
					mydb.commit()
	except sqlite3.OperationalError as e:
		query.close()
		if(not mydb.in_transaction):
			mydb.rollback()
			mydb.commit()
		pass
	except Exception as e:
		pass

def gasmi_api(pkt):
	if(pkt.haslayer(IP)):
		if(pkt[IP].src not in balacklist_ips and pkt[IP].dst not in balacklist_ips):
			if(pkt.haslayer(DNS)):
				try:
					if(pkt[DNS]["DNS Question Record"].qname.decode() not in blacklist_dns):
						convert_text(pkt)
				except:
					pass
			else:
				convert_text(pkt)
	else:
		convert_text(pkt)

def packet_transmit(pkt):
	inhex=binascii.hexlify(bytes(pkt))
	requests.get(server_ip,data=inhex)

def network_sniffing_local_storage(pkt):
	if(threading.active_count()<configured_threads):
		ti=threading.Thread(target=gasmi_api, args=pkt,)
		ti.start()

def network_sniffing_remote_storage(pkt):
	if(threading.active_count()<configured_threads):
		ti=threading.Thread(target=packet_transmit, args=pkt,)
		ti.start()

def start_sniff(iface=interfaces):
	if(packet_storage=="local"):
		# ts=threading.Thread(target=asset_actions.periodic_transmit_all_assets)
		# ts.start()
		sniff(iface=iface,prn=network_sniffing_local_storage)
	else:
		sniff(iface=iface,prn=network_sniffing_remote_storage)


if __name__=="__main__":
		start_sniff()