MassTraceroute/masstraceroute/__init__.py

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

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()
		logging.info("Creating listener socket")
		self.listener = self.init_listener()
		self.sources = dict()
		self.last_send = 0

	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:
				pass
			logging.debug("Got data from ip = {}".format(addr))
			logging.debug(data)
			logging.debug(data[48:49])
			srcport = int.from_bytes(data[48:49], byteorder="big")
			dstport = int.from_bytes(data[50:51], byteorder="big")
			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
			else:
				logging.error("Received package with srcport or dstport non existsant?")

	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()
		srcport = getrandbits(16)
		while(srcport not in self.sources):
			self.sources[srcport] = dict()
			self.sources[srcport]["ip"] = dstip
			srcport = getrandbits(16)
	   	
		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(dstport not in self.sources[srcport]):
				self.sources[srcport][dstport] = ttl
				dstport = getrandbits(16)
			logging.debug("Sending packet srcport = {}, dstport = {}, ttl = {}, dstip = {}".format(srcport, dstport, ttl, dstip))
			sender.sendto(b'', (dstip, dstport))
			self.last_send = time()

		sender.close()