pyvtraceroute/tracer.py

# -*- coding: utf-8 -*-#

# Copyright 2015 Weber Yann
#
# This file is part of pyMapTraceroute.
#
#        pyMapTraceroute is free software: you can redistribute it and/or modify
#        it under the terms of the GNU General Public License as published by
#        the Free Software Foundation, either version 3 of the License, or
#        (at your option) any later version.
#
#        pyMapTraceroute is distributed in the hope that it will be useful,
#        but WITHOUT ANY WARRANTY; without even the implied warranty of
#        MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#        GNU General Public License for more details.
#
#        You should have received a copy of the GNU General Public License
#        along with pyMapTraceroute.  If not, see <http://www.gnu.org/licenses/>.
#

import os, re, GeoIP, time, pygame


from subprocess import Popen, PIPE
from pygame import gfxdraw

import vtracemap
import vtracert

class Tracer:
	
	## Tracer instanciation
	# @param vtr The Vtracert
	def __init__(self, vtr):
		
		#Check vtr argument
		if not isinstance(vtr, vtracert.Vtracert):
			raise TypeError('Vtracert waited')
		
		self.vtr = vtr
		
		#Init object properties
		self.initTrace()
		self.gi = GeoIP.open("/usr/share/GeoIP/GeoIPCity.dat",GeoIP.GEOIP_STANDARD)
		cproc = os.popen('curl --silent zered.org/getip.php')
		self.src = cproc.readline().strip()
		
		
		self.prev = None
		self.prevLonLat = None
		
		self.selected = [] #selected hop
		
		self.regetip = re.compile('^.*\(([^\)]*)\).*$');
		self.reghost = re.compile('^ *[0-9]* *([^\(]*)\(.*$');
		self.regHop = re.compile('^[^0-9]*([0-9]*).*$');
		
		pass
	
	## Object trace initialisation
	def initTrace(self):
		self.ips = []
		self.hosts = []
		self.pts = []
		self.hops = []
		self.coords = []
		self.selected = []
	
	def progressStatus(self):
		self.vtr.drawStatus("Traceroute in progress : ...", self.vtr.font['status'],self.vtr.col['trace_progress'])
		pass
		
	
	## Generator that yield ips from a running traceroute
	# @param The host to traceroute to
	# @yield Hop's IP
	def traceGen(self,dest_name):
		
		#Init tracer
		self.initTrace();
		
		#Store source ip as first hop
		self.storeIp(self.src)
		self.hosts.append('localhost')
		self.hops.append('0')
		
		#Update status bar
		self.progressStatus()
		self.vtr.flip()
		
		#Run traceroute process
		p = Popen(['traceroute', dest_name], stdout=PIPE)
		
		#First line give us destination ip
		line = p.stdout.readline()
		
		ml = self.regetip.match(line)
		if ml != None:
			dest_ip = ml.group(1)
			
			(dlon,dlat) = self.ip2coord(dest_ip)
			(x,y) = self.vtr.trmap.latlon2ortho(dlon, dlat)

			pygame.gfxdraw.circle(self.vtr.surf['trace'],x,y,3,self.vtr.col['trace_end'])
			
		else:
			dest_ip = None
			
		self.vtr.flip()
		yield self.src
		
		#Reading traceroute output
		while True:
			line = p.stdout.readline()
			if not line:
				break
			
			print line
			
			#Getting hop number
			ml = self.regHop.match(line)
			if ml != None:
				self.hops.append(ml.group(1))
			else:
				self.hops.append('?')
			
			#hop hostname
			ml = self.reghost.match(line)
			if ml != None:
				self.hosts.append(ml.group(1))
			else:
				self.hosts.append('?')
			#hop ip
			ml = self.regetip.match(line)
			if ml != None:
				ip = ml.group(1)
			else:
				ip = None
			
			self.storeIp(ip)
			
				
			strStatus = "Traceroute in progress : "
			strStatus += self.hops[-1]+' '+self.hosts[-1]
			if ip != None:
				strStatus += '('+ip+')'
			
			self.vtr.drawStatus(strStatus, self.vtr.font['status'],self.vtr.col['trace_progress'])	
			
			self.drawRoute()
			yield self.ips[-1]
		
		if dest_ip != None:
			self.storeIp(dest_ip)
			self.hops.append('Dest')
			self.hosts.append(dest_name)
			yield dest_ip
	
	## Store an hop's ip in object properties
	# @param ip The ip adress to store
	def storeIp(self,ip):
		
		nopts = False
		
		self.ips.append(ip)
		
		if ip != None:
			lonlat = self.ip2coord(ip)
			if lonlat != None:
				#print gir['latitude'], gir['longitude']
				self.pts.append(lonlat)
			else:
				nopts = True
		else:
			nopts = True
		
		if nopts:
			self.pts.append(None)
		pass
	
	## Get geoip info for ip
	# @param ip to localisate
	def ip2coord(self,ip):
		gir = self.gi.record_by_addr(ip)
		if gir != None:
			return ( gir['longitude'], gir['latitude'])
		return None
	
	## Populate data by running a traceroute
	# @param dest_name The host
	def trace(self,dest_name):
		#Traceroute generator
		tr = self.traceGen(dest_name)
		for ip in tr:
			if ip != None:
				pass #self.storeIp(ip)
	
	## Run a traceroute and interactively update the screen
	# @param dest_name The host
	def liveTrace(self,dest_name):
		
		#Traceroute generator
		tr = self.traceGen(dest_name)
		self.initDraw()

		for ip in tr:
			if ip != None:
				#New hop
				pts = self.pts[-1]
				if pts != None:
					dcoord = self.drawNext(pts)
					self.coords.append(dcoord)
				else:
					self.coords.append(None)
			#Update the screen
			self.vtr.flip()
		pass
	
	## Init a new trace drawing
	def initDraw(self):
		self.prev = None
		self.coords = []
		self.vtr.surf['trace'].fill((0,0,0,0))
		pass
	
	## Draw the next hop in a route on the map
	# Hops are circle (filled when selected) and
	# the route is represented by line between circles
	# @param (lon,lat) Hop coordinates
	# @param selected A boolean
	# @return The coordinates of the hop on the screen as (x,y)
	def drawNext(self, (lon, lat), selected = False, last = False):
		
		if not selected :
			#If not selected draw a circle
			cdrawfun = pygame.gfxdraw.circle
		else:
			#If selected draw a filled circle
			cdrawfun = pygame.gfxdraw.filled_circle
		
		(x,y) = self.vtr.trmap.latlon2ortho(lon, lat)
		if self.prev == None:
			#Route begining
			cdrawfun(self.vtr.surf['trace'],x,y,3,self.vtr.col['trace_init'])
		else:
			#A new hop
			
			if last:
				cdrawfun(self.vtr.surf['trace'], x,y,3,self.vtr.col['trace_hop'])
			else:
				cdrawfun(self.vtr.surf['trace'], x,y,5,self.vtr.col['trace_hop'])
			
			(px,py) = self.prev
			
			#Trying to find a symetric closer from the new hop
			dw = self.vtr.trmap.scaledWidth
			spx1 = px - dw # x prev symetric 1
			spx2 = px + dw # x prev symetric 2
			
			(plon,_) = self.prevLonLat

			splon1 = plon - 360 # prev lon symetric 1
			splon2 = plon + 360 # prev lon symetric 2
			
			sd1 = abs(lon-splon1)
			sd2 = abs(lon-splon2)
			
			
			if sd1 < sd2:
				spx = spx1
				sx = x + dw
				sd = sd1
			else:
				spx = spx2
				sx = x - dw
				sd = sd2
			
			d = abs(lon-plon)
			
			#drawLine = Tracer.drawArcLine
			drawLine = pygame.gfxdraw.line
			
			if sd > d:
				drawLine(self.vtr.surf['trace'], px,py,x,y, self.vtr.col['trace_trace'])
			else:
				#symetric is closer
				drawLine(self.vtr.surf['trace'], spx,py,x,y, self.vtr.col['trace_trace'])
				drawLine(self.vtr.surf['trace'], px,py,sx,y, self.vtr.col['trace_trace'])
		
		self.prevLonLat = (lon,lat)
		self.prev = (x,y)
		
		return (x,y)
	
	## Draw an existing route on the map
	# @param delay The time to wait between the draw of two hops (in float, seconds)
	def drawTrace(self, delay = 0.1):
		
		self.initDraw()
		
		cnt = 1
		for i,pt in enumerate(self.pts):
			
			if pt != None:
				dcoord = self.drawNext(pt, (i in self.selected), (i == len(self.pts)-1) )
				self.coords.append(dcoord)
				
				if delay > 0:
					self.drawRoute(cnt)
					self.vtr.flip()
					time.sleep(delay)
				
			else:
				self.coords.append(None)
			cnt += 1
		self.drawRoute()
		pass
			
			
	## Draw the text representation of an existing route
	# One line per hop at the bottom left of the screen
	# @param stopAt Print the last 'stopAt' hops
	def drawRoute(self, stopAt = 32):
		
		strStatus=''
		
		self.vtr.surf['tracetxt'].fill((0,0,0,0))
		lh = self.vtr.font['trace_txt'].get_linesize()
		
		cy = self.vtr.surf['tracetxt'].get_height()-4
		cy -= lh
		
		if stopAt < 0 or stopAt > len(self.ips):
			rng = range(len(self.ips))
		else:
			rng = range(stopAt)
		rng.reverse()
		for i in rng:
			strStatus = self.hops[i]+' '+self.hosts[i]
			if self.ips[i] != None:
				strStatus += '('+self.ips[i]+')'
			
			if i in self.selected:
				col = self.vtr.col['trace_txtsel'] #selected color
			else:
				col = self.vtr.col['trace_txt'] #normal color
				
			statusSurf = self.vtr.font['trace_txt'].render(strStatus, 1, col)
			
			self.vtr.surf['tracetxt'].blit(statusSurf, (4,cy))
			cy -= lh
		pass
	
	## Search if there is a hop near a screen coordinate
	# Will return mutliple host if the distance is equal
	# @param (x,y) The on screen coordinates
	# @param maxDist The maximum distance bellow wich a hop is selected
	# @return A list of hop index
	def searchHopByCoord(self, (x,y), maxDist = 10):
		
		maxDist = maxDist **2
		
		curRes = []
		curMin = maxDist+1
		
		for (i,pos) in enumerate(self.coords):
			if pos != None:
				(cx,cy) = pos
				d = ((cx-x)**2) + ((cy-y)**2)
				if d <= curMin:
					if d == curMin:
						curRes.append(i)
					else:
						curRes = [i]
					curMin = d
					
		
		return curRes
	
	## Select hops given screen coordinate
	# @param pos On screen coordinate in a (x,y) tuple
	# @return A boolean set to True if hops were selected
	def select(self, pos):
		self.selected = self.searchHopByCoord(pos)
		return (len(self.selected)>0)
	
	## Return a string representing selected hops ip
	# @return One ip on each line, preffixed by \t
	def selectedIp(self):
		res = ""
		for i in self.selected:
			res += '\t'+str(self.hosts[i])+'('+str(self.ips[i])+')\n'
		return res