Code cleanup, made more consistent. Still sucks, though

scapy-tracehost.py

@@ -1,27 +1,37 @@
 #!/usr/bin/env python
 
 from scapy import all as scapy
+from optparse import OptionParser
 import sys 
 
+
 from os import geteuid
 
-def usage():
-	print "Usage: " + sys.argv[0] + " host"
-	sys.exit(1)
+def main():
+    usage = "%prog [options]"
+    parser = OptionParser(usage)
+    parser.add_option("-H","--host",dest="dsthost",help="Specify target host")
+    parser.add_option("-p","--port",dest="dstport",help="Specify traceroute port", default="80")
+    parser.add_option("-f","--file",dest="dstfile",help="Specify output file")
+    parser.add_option("-r","--resolve",dest="resolve",action="store_true",help="Enable AS resolution",default=False)
+
+    (options, args) = parser.parse_args()
+
+    dsthost = options.dsthost
+    dstport = options.dstport
+    dstfile = options.dstfile
+    if options.resolve == False:
+        scapy.conf.AS_resolver=None
 
-def tracehost(host):
-	res,unans = scapy.traceroute(host)
-	res.graph(type="svg",target=">"+host+"-traceroute.svg")
+    tracehost(dsthost,dstport,dstfile)
 
-def getargs():
-	argc = len(sys.argv)
-	if argc != 2:
-		usage()
 
-	tracehost(sys.argv[1])
+def tracehost(dsthost,dstport,dstfile):
+	res,unans = scapy.traceroute(dsthost, dport=int(dstport))
+	res.graph(type="svg",target=">"+dstfile)
 
 if __name__ == '__main__':
-	if geteuid() !=0:
-		print "This script must be run as root. Sorry..."
-		sys.exit(1)
-	getargs()
+    if geteuid() !=0:
+        print "[-] This script must be run as root. Sorry..."
+        sys.exit(1)
+    main()

scapy-tracenet.py

@@ -1,43 +1,50 @@
 #!/usr/bin/env python
 
 from scapy import all as scapy
-import sys 
-
+from optparse import OptionParser
+import sys
 from mmap import mmap, MAP_PRIVATE, PROT_READ
 from os import fstat
 from os import geteuid
 
-def usage():
-	print "Usage: " + sys.argv[0] + " customer hostfile"
-	sys.exit(1)
+def main():
+    usage = "%prog [options]"
+    parser = OptionParser(usage)
+    parser.add_option("-H","--hosts",dest="dsthostsfile",help="Specify hosts file")
+    parser.add_option("-p","--port",dest="dstport",help="Specify traceroute port", default="80")
+    parser.add_option("-f","--file",dest="dstfile",help="")
+    parser.add_option("-r","--resolve",dest="resolve",action="store_true",help="Enable AS resolution",default=False)
+
+    (options, args) = parser.parse_args()
+
+    dsthostsfile = options.dsthostsfile
+    dstport = options.dstport
+    dstfile = options.dstfile
+    if options.resolve == False:
+        scapy.conf.AS_resolver=None
 
-def tracehost(customer, thostsfile):
-	f = file(thostsfile,mode='rt')
+    tracehost(dsthostsfile,dstport,dstfile)
+
+def tracehost(dsthostsfile,dstport,dstfile):
+	f = file(dsthostsfile,mode='rt')
 	fd = f.fileno()
 	m = mmap(fd, fstat(fd).st_size, MAP_PRIVATE, PROT_READ)
 
-	thosts=[]
+	dsthosts=[]
 	while True:
 	        line = m.readline()
         	if not line: break
-        	thosts.extend(line.split())
+        	dsthosts.extend(line.split())
 
-	res,unans = scapy.traceroute(thosts)
-	res.graph(type="svg",target=">" +customer+"-traceroute.svg")
+	res,unans = scapy.traceroute(dsthosts,dport=int(dstport))
+	res.graph(type="svg",target=">"+dstfile)
 
 #	for thost in thosts:
 #        	res1,unans1 = scapy.traceroute(thost)
 #        	res1.graph(type="ps",target=">/tmp/"+cust+"-traceroute-"+thost)
 
-def getargs():
-	argc = len(sys.argv)
-	if argc != 3:
-		usage()
-
-	tracehost(sys.argv[1],sys.argv[2])
-
-if __name__ == '__main__':
-	if geteuid() !=0:
-		print "This script must be run as root. Sorry..."
-		sys.exit(1)
-	getargs()
+if __name__=='__main__':
+    if geteuid() !=0:
+        print "[-] This script must be run as root. Sorry..."
+        sys.exit(1)
+    main()