Minor enhancement to internal --profile function

lib/core/common.py

@@ -1093,50 +1093,74 @@ def isBase64EncodedString(subject):
 def isHexEncodedString(subject):
     return re.match(r"\A[0-9a-fA-F]+\Z", subject) is not None
 
-def profile(profileOutputFile=None, imageOutputFile=None):
+def profile(profileOutputFile=None, dotOutputFile=None, imageOutputFile=None):
+    try:
+        from extra.gprof2dot import gprof2dot
+        from extra.xdot import xdot
+        import gobject
+        import gtk
+        import pydot
+    except ImportError, e:
+        logger.error(e)
+        return
+
     if profileOutputFile is None:
         profileOutputFile = os.path.join(paths.SQLMAP_OUTPUT_PATH, "sqlmap_profile.raw")
 
+    if dotOutputFile is None:
+        dotOutputFile = os.path.join(paths.SQLMAP_OUTPUT_PATH, "sqlmap_profile.dot")
+
     if imageOutputFile is None:
         imageOutputFile = os.path.join(paths.SQLMAP_OUTPUT_PATH, "sqlmap_profile.png")
 
     if os.path.exists(profileOutputFile):
         os.remove(profileOutputFile)
 
+    if os.path.exists(dotOutputFile):
+        os.remove(dotOutputFile)
+
     if os.path.exists(imageOutputFile):
         os.remove(imageOutputFile)
 
     infoMsg  = "profiling the execution into file %s" % profileOutputFile
     logger.info(infoMsg)
 
+    # Start sqlmap main function and generate a raw profile file
     cProfile.run("start()", profileOutputFile)
 
-    infoMsg  = "converting profile data into a graph image '%s'" % imageOutputFile
+    infoMsg = "converting profile data into a dot file '%s'" % dotOutputFile
     logger.info(infoMsg)
 
-    graphScriptPath = os.path.join(paths.SQLMAP_EXTRAS_PATH, 'gprof2dot', 'gprof2dot.py')
+    # Create dot file by using extra/gprof2dot/gprof2dot.py
+    # http://code.google.com/p/jrfonseca/wiki/Gprof2Dot
+    dotFilePointer = open(dotOutputFile, 'wt')
+    parser = gprof2dot.PstatsParser(profileOutputFile)
+    profile = parser.parse()
+    profile.prune(0.5/100.0, 0.1/100.0)
+    dot = gprof2dot.DotWriter(dotFilePointer)
+    dot.graph(profile, gprof2dot.TEMPERATURE_COLORMAP)
+    dotFilePointer.close()
 
-    # TODO: find the Windows version of Unix command 'dot'
-    process = execute('%s %s -f pstats %s | dot -Tpng -o %s' % (sys.executable, graphScriptPath, profileOutputFile, imageOutputFile), shell=True, stdout=None, stderr=PIPE)
-    processStderr = process.communicate()[1]
+    infoMsg = "converting dot file into a graph image '%s'" % imageOutputFile
+    logger.info(infoMsg)
 
-    if processStderr or not os.path.exists(imageOutputFile):
-        errMsg = "there was an error while converting ('%s'), " % processStderr.strip()
-        errMsg += "but you can still find raw profile data "
-        errMsg += "inside file '%s'" % profileOutputFile
-        logger.error(errMsg)
+    # Create graph image (png) by using pydot (python-pydot)
+    # http://code.google.com/p/pydot/
+    pydotGraph = pydot.graph_from_dot_file(dotOutputFile)
+    pydotGraph.write_png(imageOutputFile)
 
-        return
+    infoMsg = "displaying interactive graph with xdot library"
+    logger.info(infoMsg)
+
+    # Display interactive Graphviz dot file by using extra/xdot/xdot.py
+    # http://code.google.com/p/jrfonseca/wiki/XDot
+    win = xdot.DotWindow()
+    win.connect('destroy', gtk.main_quit)
+    win.set_filter("dot")
+    win.open_file(dotOutputFile)
+    gobject.timeout_add(1000, win.update, dotOutputFile)
+    gtk.main()
 
-    try:
-        if PLATFORM == 'mac':
-            subprocess.call(('open', imageOutputFile))
-        elif PLATFORM == 'posix':
-            subprocess.call(('xdg-open', imageOutputFile))
-        elif PLATFORM == 'nt':
-            subprocess.call(('start', imageOutputFile))
-    except:
-        pass
 
 def getConsoleWidth(default=80):
     width = None