pyvtraceroute/vtracert.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/>.
#
control_summary = [
"+----------------------+----------------------------------+",
"|                Controls summary                         |",
"+----------------------+----------------------------------+",
"+----------------------+----------------------------------+",
"| ^h, ?                |  display this help               |",
"+----------------------+----------------------------------+",
"| esc, ^q              |  exit the programm or this help  |",
"+----------------------+----------------------------------+",
"+----------------------+----------------------------------+",
"| [a-zA-Z0-9\.]        |  Host input                      |",
"+----------------------+----------------------------------+",
"| mid-mouse btn        |  paste clipboard into host input |",
"+----------------------+----------------------------------+",
"| enter                |  validate host input and         |",
"|                      |  run traceroute                  |",
"+----------------------+----------------------------------+",
"| ^c                   |  clear host input                |",
"+----------------------+----------------------------------+",
"+----------------------+----------------------------------+",
"| mouse scroll, +, -   : zoom in/out                      |",
"+----------------------+----------------------------------+",
"| arrow keys           : move on the map                  |",
"+----------------------+----------------------------------+",
"| left mouse btn press : select hops on the map           |",
"+----------------------+----------------------------------+"
]

import random
import time
import os

import pygame
from pygame import gfxdraw

import tracer
import vtracemap

class Vtracert:
	
	## Instanciate a visual tracert
	# @param mapdatfile The filename containing the borders coordinates
	# @param width Window width
	# @param height Window height
	# @param cols Override default colors, its a dict of tuple (r,g,b) (r,g,b,a), valid keys are : bg, map, prompt, status, trace_init, trace_hop, trace_trace
	# @param fonts Override default font type and size its a dict of tuple (fontType, fontSize), valid keys are : prompt, status
	# @param promptText The prompt when reading the host to traceroute
	# @param defaultInput Default host to traceroute
	def __init__(self, mapdatfile, width = 1024, height = 768, cols = dict(), fonts = dict(), promptText = "Enter an ip or an hostname and press enter to traceroute :    %s_", defaultInput = "gnu.org"):
		
		pygame.init() #Pygame init
		
		#Options
		self.mapfile = mapdatfile
		self.width = width
		self.height = height
		#Colors
		self.col = dict()
		self.col['alpha'] = pygame.Color(0,0,0,0)
		self.col['bg'] = pygame.Color(0,0,0)
		self.col['map'] = pygame.Color(15,8,63)
		self.col['prompt'] = pygame.Color(200,200,200)
		self.col['status'] = pygame.Color(200,200,200)
		self.col['cur_status'] = self.col['status']
		self.col['cur_statusline'] = self.col['status']
		self.col['trace_init'] = pygame.Color(0,0xf0,0)
		self.col['trace_end'] = pygame.Color(0xee,0x33,0xee)
		self.col['trace_hop'] = pygame.Color(0xf0,0xf0,0)
		self.col['trace_trace'] = pygame.Color(0xf0,0,0)
		self.col['trace_txt'] = pygame.Color(200, 200, 200, 125)
		self.col['trace_txtsel'] = self.col['trace_hop']
		self.col['trace_progress'] = pygame.Color(0,0xf0,0)
		self.col['help_txt'] =  pygame.Color(200,200,200)
		self.col['help_title'] = pygame.Color(120,30,230)
		self.col['help_bg'] = pygame.Color(0,0,0)
		for n, c in cols:
			self.col[n] = pygame.Color(c)
		#Fonts
		self.font = dict()
		self.font['prompt'] = pygame.font.Font(None, 20)
		self.font['status'] = pygame.font.Font(None,20)
		self.font['cur_status'] = self.col['status']
		self.font['trace_txt'] = pygame.font.Font(None,20)
		
		self.font['help_title'] = pygame.font.Font(None,48)
		self.font['help_title'].set_underline(True);
		self.font['help_title'].set_bold(True);
		
		self.font['help'] = pygame.font.SysFont('monospace', 16)
		for n, (ft,fs) in fonts:
			self.font[n] = pygame.font.Font(ft,fs)
		#Prompt options
		self.promptText = promptText
		self.inputText = defaultInput
		self.inputText = '1.202.15.14'
		
		#Event filtering
		pygame.event.set_allowed(None)
		pygame.event.set_allowed([pygame.MOUSEBUTTONUP,pygame.KEYDOWN, pygame.KEYUP, pygame.QUIT])

		#Surface creation
		self.surf = dict()
		self.surf['screen'] = pygame.display.set_mode((width, height), pygame.DOUBLEBUF )
		self.surf['map'] = pygame.Surface((width, height),pygame.HWSURFACE)
		self.surf['trace'] = pygame.Surface((width, height),pygame.HWSURFACE | pygame.SRCALPHA)
		self.surf['tracetxt'] = pygame.Surface((width, height),pygame.HWSURFACE | pygame.SRCALPHA)
		self.surf['status'] = None
		
		for s in self.surf:
			if self.surf[s] != None:
				self.surf[s].fill(self.col['alpha'])
		
		self.help(0)
		self.drawStatus('Loading...');
		pygame.display.flip()
		
		#Clipboard handling module initialisation
		pygame.scrap.init()
		
		#Keyboard vars
		self.caps = False
		self.ctrl = False
		
		#Map info
		self.trmap = vtracemap.TraceMap(self)
		self.tracer = tracer.Tracer(self)
		
	
	## Color accessor
	# @param cname Color name
	# @param v The value to set or False to get
	# @return The color named cname
	# @see Vtracert::opt()
	def Col(self, cname, v = False):
		return Vtracert.opt(cname, self.col, v)
	## Surface accessor
	# @param sname surface name
	# @param v The value to set or False to get
	# @return The surface named sname
	# @see Vtracert::opt()
	def Surf(self, sname, v = False):
		return Vtracert.opt(sname, self.surf, v)
	## Font accessor
	# @param fname font name
	# @param v The value to set or False to get
	# @return The font named fname
	# @see Vtracert::opt()
	def Font(self, fname, v = False):
		return Vtracert.opt(sname, self.font, v)
	
	## A generic class to handle dict properties
	# @param cname The dict key
	# @param dopt The dict
	# @param v The value or False for get and not set
	# @return dopt[cname]
	# @see Vtracert::Col(), Vtracert::Surf(), Vtracert::Font()
	@classmethod
	def opt(c, cname, dopt, v):
		if not (not v and isinstance(v,bool)):
			dopt[cname] = v
		return dopt[cname]
	
	## Return the unicode representation of the combiation of ctrl+a letter key
	@classmethod
	def ctrlUnicode(classo, c):
		return unicode(chr(ord(c.lower())-ord('a')+1))
		pass
	
	##Display an help on the screen and wait for an KEYDOWN or QUIT event
	def help(self, ttl = None):
		scr = self.surf['screen']
		helpTxts = [ 
			"Vtraceroute help",
			"",
			"Vtraceroute is a traceroute program ( more info : man traceroute )",
			"that displays hops on a map using geoipLookup.",
			"",
			"Controls summary :",
			"",
			"^h, ?                : display this help",
			"esc, ^q              : exit the programm or this help",
			"",
			"[a-zA-Z0-9\.]        : Host input",
			"mid-mouse btn        : paste clipboard into host input",
			"enter                : validate host input and",
			"                       run traceroute",
			"^c                   : clear host input",
			"",
			"",
			"mouse scroll, +, -   : zoom in/out",
			"arrow keys           : move on the map",
			"left mouse btn press : select hops on the map",
			"","",
			"All borders data extracted from data provided by http://thematicmapping.org and",
			"Licenced under Creatice Commons Attribution-Share Alike License 3.0"
			"",
			"Press any key or mouse button to exit this help"
		]
		
		tFont = self.font['help_title']
		hFont = self.font['help']
		hCol = self.col['help_txt']
		tCol = self.col['help_title']
		
		ty = 50
		lspacing = 3
		hy = ty + tFont.get_linesize() + lspacing
		hx = 40
		
		#Fill screen
		scr.fill(self.col['help_bg'])
		
		#Print title
		titleTxt = helpTxts.pop(0)
		tsurf = tFont.render(titleTxt,1,tCol)
		
		tsc = tsurf.get_rect()
		tsc.centerx = self.width/2 #center the title
		tsc.y = ty
		
		scr.blit(tsurf,tsc)
		
		#Print help text
		for line in helpTxts:
			
			if len(line) > 0:
				tsurf = hFont.render(line,1,hCol)
				scr.blit(tsurf,(hx, hy))
			
			hy += hFont.get_linesize()
			hy += lspacing
			#pygame.display.flip();time.sleep(0.05);
		
		pygame.display.flip()
		
		if ttl == None:
			wait = True
			while wait:
				evts = pygame.event.get()
				for evt in evts:
					if evt.type in (pygame.QUIT, pygame.KEYDOWN, pygame.MOUSEBUTTONUP):
						wait = False
						break;
		else:
			time.sleep(ttl)
		pass
	
	"""
	## Return True 
	# @c1 unicode1
	# @c2 unicode2
	def isCtrlUnicode(self, c, cu):
		return (self.ctrl and Vtracert.ctrlUnicode(c) == cu)
	"""
	
	## Blit a surface on the screen surface
	# @param sname The surface name
	# @param pos A tuple (x,y) for the blit position
	# @return self.surf['screen'].blit(self.surf[sname], pos)
	def blitScreen(self, sname, pos = (0,0) ):
		return self.surf['screen'].blit(self.surf[sname], pos)
	
	## Draw the map again and call Vtracert::flip()
	# Call Vtracert::fillScreen()
	# @see Vtracert::fillScreen(), Vtracert::flip()
	def mapRedraw(self):
		self.fillScreen()
		self.trmap.draw()
		self.flip()
	
	## Fill the screen with the 'bg' color
	def fillScreen(self):
		self.surf['screen'].fill(self.col['bg'])
		pass
	
	## Blit all the surface is the right order but don't call pygame.display.flip()
	# The status bar horizontal ruler is a draw and not a blit
	# @see Vtracert::flip()
	def blit(self):
		self.blitScreen('map')
		self.blitScreen('status')
		self.drawStatusLine()
		self.blitScreen('trace')
		self.blitScreen('tracetxt')
		pass
	
	## Blit all surface and flip display
	# @see Vtracert::blit()
	def flip(self, fill = True):
		if fill:
			self.fillScreen()
		self.blit()
		pygame.display.flip()
		pass
	
	## Draw the prompt in the status surface
	# @see Vtracert::drawStatus()
	def drawPrompt(self):
		self.drawStatus(self.promptText%(self.inputText), self.font['prompt'], self.col['prompt'])
	
	## Draw the status bar on the screen
	# @param statusText The text to display in the status bar
	# @param font The font used to render the text
	# @param col The text color
	# @param antialiasing A boolean telling if the text is rendered antialiased
	# @param line A boolean telling if the status bar line has to be drawn
	# @param lineCol The status bar line color (by default same as col)
	def drawStatus(self, statusText, font = None, col = None, antialiasing = True, line = True, lineCol = None):
		
		if font == None:
			font = self.font['status']
		self.font['cur_status'] = font
		
		if col == None:
			col = self.col['status']
		self.col['cur_status'] = col
		
		if lineCol == None:
			lineCol = col
		self.col['cur_statusline'] = lineCol
		
		self.surf['status'] = font.render(statusText, antialiasing, col)
		
		self.blitScreen('status')
		
		if line:
			self.drawStatusLine()
		
		pass
	
	## Draw the status bar horizontal ruler
	# @see Vtracert::drawStatus()
	def drawStatusLine(self):
		ly = self.font['cur_status'].get_linesize() + 2
		pygame.gfxdraw.line(self.surf['screen'], 0, ly, self.width - 1, ly, self.col['cur_statusline'])
		pass
	
	## Handle a KEYUP event
	# @param evt Event
	def keyupEvt(self, evt):
		if evt.key == 303 or evt.key == 304: #caps release
			self.caps = False
		elif evt.key == 305 or evt.key == 306: #ctrl release
			self.ctrl = False
		pass
	
	## Handles a KEYDOWN event
	# @param evt Event
	def keydownEvt(self, evt):
		ccode = evt.key
		
		c2u = Vtracert.ctrlUnicode
		
		allowedChar = range(ord('a'),ord('z')+1) 
		allowedChar += range(ord('A'),ord('Z')+1)
		allowedChar += range(ord('0'),ord('9')+1)
		allowedChar += [ord('.'),ord('\b'),ord('\r'),ord('\n')]
		
		allowedChar = [ chr(cc) for cc in allowedChar ]
		ctrlChar = [	c2u('q'), #^q
						c2u('h'), #^h
						c2u('n'), #^n
						c2u('c'), #^n
						'?',
						u'\x1b' #escape button
					]
		
		#print evt
		#print self.ctrl
		
		if ccode in [273,274,275,276]: #Arrow keys
			if ccode == 273:
				d = 'up'
			elif ccode == 274:
				d = 'down'
			elif ccode == 275:
				d = 'right'
			else:
				d = 'left'
			
			self.trmap.zoom( (None, None), d)
			self.tracer.drawTrace(0)
			self.flip()
			pygame.event.clear(pygame.MOUSEBUTTONUP)
			
		elif ccode == 303 or ccode == 304: #caps
			self.caps = True
		elif evt.key == 305 or evt.key == 306: #ctrl release
			self.ctrl = True
		elif self.ctrl or cc in ctrlChar: # ctrl + key handle
			
			cc = evt.unicode
			if cc == u'\x1b' or cc == c2u('q'): #quit (esc or ^q)
				pygame.event.post(pygame.event.Event(pygame.QUIT))
			elif cc == c2u('h') or cc == '?': #help ^h ? ^?
				print "Help !!!!!!!!"
				self.help()
				self.flip()
			elif cc == c2u('c'):
				self.inputText = ''
				self.drawPrompt()
			elif cc == c2u('n'):
				for cname in self.col:
					c = self.col[cname]
					c.r = 255 - c.r
					c.g = 255 - c.g
					c.b = 255 - c.b
					self.col[cname] = c
				self.mapRedraw()
				
		elif evt.unicode in allowedChar:
			
			cc = evt.unicode
				
			if cc == '\b': #backspace
				if len(self.inputText) > 0:
					self.inputText = self.inputText[:-1]
			elif cc == '\r' or cc == '\n':
				print "OK, tracerouting %s ..."%self.inputText
				dest = self.inputText
				self.inputText = ""
				
				self.surf['trace'].fill((0,0,0,0))
				
				self.tracer.liveTrace(dest)
				#self.tracer.trace(dest)
				self.flip()
				
			else:
				self.inputText += cc
			
			self.drawPrompt()
		pass
	
	## Handles a MOUSEBUTTONUP event
	# @param evt Event
	def mouseboutonupEvt(self, evt):
		if evt.button == 4 or evt.button == 5: #scrollzoom
				
			if evt.button == 4 or evt.button == 1:
				zoom = True
			elif evt.button == 5 or evt.button == 3:
				zoom = False
				
			self.trmap.zoom(evt.pos,zoom)
			self.tracer.drawTrace(0)
			self.blit()
			pygame.event.clear(pygame.MOUSEBUTTONUP)
		#end scrollzoom
		
		elif evt.button == 1:
			#Trying to select hops
			if self.tracer.select(evt.pos):
				print 'Selected : \n',self.tracer.selectedIp()
			self.tracer.drawTrace(0)
			
		elif evt.button == 2:
			clipTxt = pygame.scrap.get('TEXT')
			if clipTxt:
				inputTxt += clipTxt
				self.drawPrompt()
		else:
			print "\n\n"
			types = pygame.scrap.get_types()
			for t in types:
				print t,pygame.scrap.get(t)
		
		pass
	
	
	## Listen for pygame events
	def run(self):
		
		self.drawPrompt();
		self.mapRedraw();
		self.flip()
		
		#Display small onetime help text
		htext = self.font['status'].render("Type ^h or ? for help.",1, self.col['status'])
		self.surf['screen'].blit(htext, (5, self.font['status'].get_linesize()+2))
		pygame.display.flip()
		
		again = True
		
		while again:
			evt = pygame.event.wait()
			if evt != pygame.NOEVENT:
				if evt.type == pygame.QUIT:
					again = False
					break
				elif evt.type == pygame.KEYUP:
					self.keyupEvt(evt)
				elif evt.type == pygame.KEYDOWN:
					self.keydownEvt(evt)
				elif evt.type == pygame.MOUSEBUTTONUP: 
					self.mouseboutonupEvt(evt)
				
				self.drawPrompt();
				self.flip()