import threading
import socket
import logging
from random import getrandbits
from time import sleep, time
from multiprocessing.pool import ThreadPool
from multiprocessing import Lock
from threading import Thread

from pprint import pprint

class MassTraceroute:
	def __init__(self, socket_timeout, listener_timeout, threads, targets, maxhops):
		logging.info("Setting up class")
		self.wait = 0.2
		self.maxhops = maxhops + 1
		self.listener_timeout = listener_timeout
		self.socket_timeout = socket_timeout
		self.targets = targets
		self.threads = threads
		self.pool = ThreadPool(processes=self.threads)
		self.routes = dict()
		for target in self.targets:
			self.routes[target] = dict()
		logging.info("Creating listener socket")
		self.listener = self.init_listener()
		self.sources = dict()
		self.last_send = time()

	def main(self):
		logging.info("Starting listener thread")
		Thread(target=self.process_listener).start()
		logging.info("Starting sender thread")
		Thread(target=self.process_sender).start()

	def init_listener(self):
		listener = socket.socket(
			family=socket.AF_INET,
			type=socket.SOCK_RAW,
			proto=socket.IPPROTO_ICMP
		)
		listener.settimeout(self.listener_timeout)
		listener.bind(('', 0))
		return listener

	def process_listener(self):
		while((self.last_send + self.listener_timeout) > time()):
			logging.debug("Listening...")
			try:
				data, addr = self.listener.recvfrom(1024)
			except socket.error as e:
				logging.debug(e)
				pass
			srcport = int.from_bytes(data[48:50], byteorder="big")
			dstport = int.from_bytes(data[50:52], byteorder="big")
			logging.info("Got data from ip = {}, srcport = {}, dstport = {}, original target = {}, hop = {}".format(addr[0], srcport, dstport, self.sources[srcport]["ip"], self.sources[srcport][dstport]))
			if srcport in self.sources and dstport in self.sources[srcport]:
				hop = self.sources[srcport][dstport]
				ip = self.sources[srcport]["ip"]
				self.routes[ip][hop] = addr[0]
				logging.info(self.routes[ip][hop])
			else:
				logging.error("Received package with srcport or dstport non existsant?")
		pprint(self.routes)

	def process_sender(self):
		logging.info("Starting process_sender")
		#for target in self.targets:
		#	self.send(target)

		self.pool.map(self.send, self.targets)


	def init_sender(self):
		sender = socket.socket(
			family=socket.AF_INET,
			type=socket.SOCK_DGRAM,
			proto=socket.IPPROTO_UDP
		)
		return sender

	def send(self, dstip):
		logging.info("Entering send()")
		sender = self.init_sender()
		while(True):
			srcport = getrandbits(16)
			if srcport not in self.sources:
				self.sources[srcport] = dict()
				self.sources[srcport]["ip"] = dstip
				break
	   	
		logging.info("Binding sending socket")
		sender.bind(('', srcport))

		for ttl in range(1, self.maxhops):
			sender.setsockopt(socket.SOL_IP, socket.IP_TTL, ttl)
			dstport = getrandbits(16)
			while(True):
				dstport = getrandbits(16)
				if dstport not in self.sources[srcport]:
					self.sources[srcport][dstport] = ttl
					logging.debug("Saved ttl = {}, original ttl = {}, target = {}".format(self.sources[srcport][dstport], ttl, dstip))
					break
			logging.debug("Sending packet srcport = {}, dstport = {}, ttl = {}, dstip = {}".format(srcport, dstport, ttl, dstip))
			sender.sendto(b'', (dstip, dstport))
			self.last_send = time()

		sender.close()