用python写了一个跟踪路由的小东西

转载：http://www.iteye.com/topic/550804

collect.py,  这个只能在Linux下用。。。#!/usr/bin/python

因为使用的是traceroute命令

import sys
import os
import thread
import cPickle
import time
import re

from socket import inet_ntoa
from subprocess import Popen, PIPE

import dpkt

try:
  import pcap

except ImportError:
  print "Failed to import pcap library!"
  print "If you are using Debian/Ubuntu Linux, please:"
  print " root@whatever:~# apt-get install python-pypcap"
  sys.exit(1)


class TracertException(Exception):
  pass

def tracert(host):
  """
  Traceroute 
  eg: l = tracert("8.8.8.8")

  Return value: a list of routers

  """
  p = Popen(["traceroute", "-n", "-q", "2", "-I", host], stdout=PIPE, stderr=PIPE)
  p.wait()
  (o, e) = p.communicate()
  if e != '':
    e = e.replace("\n", "")
    raise TracertException(e)
  
  t = []
  
  r = re.compile(r"((\d{1,3}\.){3}\d{1,3})") # this matches a ip address
  for i in o.split("\n")[1:-1]: # prompt at head and empty line at tail
    s = r.search(i)
    if s:
      t.append(s.group())
    else:
      t.append(None)
      
  del p
  return t


def thread_func(arg):
  "thread func, trace route and store in rmap"

  global rmap, rmap_lock, taskqueue, taskqueue_lock, ending_flag
  global threadcount, threadcount_lock, my_ip, tracked, tracked_lock
  try:
    threadcount_lock.acquire()
    threadcount = threadcount + 1
    threadcount_lock.release()
    
    while True:
      while not len(taskqueue):
        time.sleep(1)
        
        if ending_flag:
          threadcount_lock.acquire()
          threadcount = threadcount - 1
          threadcount_lock.release()
          #print "Thread %d is terminated." % thread.get_ident()
          thread.exit()
    
      if ending_flag:
        threadcount_lock.acquire()
        threadcount = threadcount - 1
        threadcount_lock.release()
        #print "Thread %d is terminated." % thread.get_ident()
        thread.exit()
      
      taskqueue_lock.acquire()
      if len(taskqueue):
        ip = taskqueue.pop(0)
      else:
        taskqueue_lock.release()
        continue  
      taskqueue_lock.release()
    
      print "Thread %d got work to do: %s" % (thread.get_ident(), ip) 
                                              
      route = tracert(ip)
      tracked_lock.acquire()
      for i in route:
        tracked[i] = 1 # None is ok, i think, so not filtering
      tracked_lock.release()
      
      
      print "%s: %s" % (ip, repr(route))
      
      f = my_ip
      h = 1
      rmap_lock.acquire()
      for t in route:
        if t == None:
          h = h + 1
          continue
        rmap[(f,t)] = h
        h = 1
        f = t
      
      if h != 1:
        rmap[(f,None)] = ip # missing
      
      rmap_lock.release()
      
  except Exception, e:
    import traceback
    threadcount_lock.acquire()
    threadcount = threadcount - 1
    threadcount_lock.release()
    print "Thread %d is *CRASHED*!" % (thread.get_ident(), )
    print traceback.format_exc()
    thread.exit()

def get_ifaddr(ifn):
  """get ip address of specified iface
     >>> get_ifaddr("wlan0")
     "113.121.xx.xx"
  """
  import struct, socket, fcntl
  s = socket.socket()
  return socket.inet_ntoa(fcntl.ioctl(s.fileno(), 0x8915, struct.pack("256s", ifn[:15]))[20:24])  

def main():
  global rmap, rmap_lock, taskqueue, taskqueue_lock, ending_flag
  global threadcount, threadcount_lock, tracked, tracked_lock, my_ip
  
  if len(sys.argv) < 2:
    print "Please specify a network interface!"
    sys.exit(1)

  
  if os.geteuid():
    print "This tool require root privilege to run!"
    sys.exit(1)
  
  if os.path.exists("routemap.dat"): # is there a saved router map?
    print "Loading saved route map..."
    f = open("routemap.dat", "rb")
    rmap, my_ip = cPickle.load(f)
    f.close()
    
  my_ip = get_ifaddr(sys.argv[1])
  
  # add my ip to tracked
  tracked[my_ip] = 1
  
  # start working threads
  for i in xrange(20):
    thread.start_new_thread(thread_func, (None,))
  
  # start packet capture
  try:
    p = pcap.pcap(sys.argv[1])
    
    for i in p:
      pkt = dpkt.ethernet.Ethernet(i[1])
      while not isinstance(pkt, str):
        pkt = pkt.data
        if isinstance(pkt, dpkt.ip.IP):
          src = inet_ntoa(pkt.src)
          dst = inet_ntoa(pkt.dst)
          #print "%s => %s" % (src, dst)
          if not tracked.get(src):
            tracked[src] = 1
            print "Task: %s" % src
            taskqueue_lock.acquire()
            taskqueue.append(src)
            taskqueue_lock.release()
            
          
          if not tracked.get(dst):
            tracked[dst] = 1
            print "Task: %s" % dst
            taskqueue_lock.acquire()
            taskqueue.append(dst)
            taskqueue_lock.release()
          
  except KeyboardInterrupt:
    intr = True
    if taskqueue_lock.locked():
      taskqueue_lock.release()
    
  if intr:
    print "Waiting for all thread to terminate..."
    try:
      ending_flag = True
      while threadcount:
        print threadcount
        time.sleep(1)
    except KeyboardInterrupt:
      print "OK, OK, I'll terminate immediatly..."
      sys.exit(1)
      
  
  # save the result
  print "Saving route map..."
  f = open("routemap.dat", "wb")
  cPickle.dump((rmap, my_ip), f, 2)
  f.close()
  print "Program terminated."
  sys.exit(0)


class DbgLock:
  def __init__(self, name):
    self.lock = thread.allocate_lock()
    self.name = name
  
  def acquire(self):
    self.lock.acquire()
    print ":: Thread %d Locking %s" % (thread.get_ident(), self.name)
  
  def release(self):
    print ":: Thread %d Releasing %s" % (thread.get_ident(), self.name)
    self.lock.release()

  def locked(self):
    return self.lock.locked()
# ----------------------------------------------------

rmap = {} # Router map, all things goes here. this should be saved.
#rmap_lock = thread.allocate_lock()
rmap_lock = DbgLock("rmap")

taskqueue = [] # task queue, ips going to be tracked
#taskqueue_lock = thread.allocate_lock()
taskqueue_lock = DbgLock("taskqueue")

tracked = {} # ips already tracerouted
#tracked_lock = thread.allocate_lock()
tracked_lock = DbgLock("tracked")

ending_flag = False # whether to end collecting...

threadcount = 0
#threadcount_lock = thread.allocate_lock()
threadcount_lock = DbgLock("threadcount")

my_ip = ""

if __name__ == "__main__":
  main()

　　drawgraph_graphviz.py

#!/usr/bin/python

import sys
import cPickle
import yapgvb
import random

def main():
  
  def get_vertex_info(ip):
    if nodetree.get(ip): # is source ip node already exist?
      vv = nodetree[ip][0] # vertex
      vl = nodetree[ip][1] # child vertex list
     
    else:
      vv = g.add_node(ip)
      #vv.shape = "circle"
      vv.color = "blue"
      vl = []
      nodetree[ip] = (vv, vl)
    
    return (vv, vl)
  
  nodetree = {}
  
  filename = "routemap.dat"
  if len(sys.argv) < 2:
    print "Didn't specify filename, use 'routemap.dat'"
  else:
    filename = sys.argv[1]
  
  f = open(filename, "rb")
  rmap, my_ip = cPickle.load(f)
  f.close()
  
  g = yapgvb.Digraph("RouteMap")
  
  get_vertex_info(my_ip)[0].color = "green"
  
  for v, deg in rmap.items():
    
    vertex, vertexlist = get_vertex_info(v[0])
    dup = False
    for ip, data in vertexlist:
      if ip == v[1]: # this is a duplicate
        dup = True
        break
    
    if dup:
      continue
    
    if v[1] is not None: # is it a LOST trace?
      if deg > 1:
        for i in range(1, deg):
          vt = g.add_node("* %d" % random.randint(0,100000000))
          vt.color = "red"
          vertex >> vt
          vertex = vt
          
      vertex_t, vertex_tlist = get_vertex_info(v[1])
      if vertex == vertex_t:
        print v
      else:
        edge = vertex >> vertex_t
        vertexlist.append((v[1], edge))
        
        
    else:

      vertex_t = g.add_node("LOST(%s)" % deg)
      vertex_t.color = "red"
      if vertex == vertex_t:
        print v
      else:
        edge = vertex >> vertex_t
  
  g.layout(yapgvb.engines.dot)
  
  print "Rendering..."
  g.render("result.svg")

if __name__ == "__main__":
  main()