Minor code refactoring

lib/core/common.py

@@ -22,6 +22,7 @@ with sqlmap; if not, write to the Free Software Foundation, Inc., 51
 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 """
 
+import cProfile
 import os
 import random
 import re
@@ -54,6 +55,7 @@ from lib.core.exception import sqlmapMissingDependence
 from lib.core.exception import sqlmapSyntaxException
 from lib.core.settings import DESCRIPTION
 from lib.core.settings import IS_WIN
+from lib.core.settings import PLATFORM
 from lib.core.settings import SITE
 from lib.core.settings import SQL_STATEMENTS
 from lib.core.settings import SUPPORTED_DBMS
@@ -1087,44 +1089,47 @@ def isBase64EncodedString(subject):
 def isHexEncodedString(subject):
     return re.match(r"\A[0-9a-fA-F]+\Z", subject) is not None
 
-def profile(profileOutputFile='sqlmap.profile', imageOutputFile='profile.png'):
-    import cProfile
-    cProfile.run("start()", profileOutputFile)
+def profile(profileOutputFile=None, imageOutputFile=None):
+    if profileOutputFile is None:
+        profileOutputFile = os.path.join(paths.SQLMAP_OUTPUT_PATH, "sqlmap_profile.raw")
 
-    graphScriptPath = os.path.join(paths.SQLMAP_EXTRAS_PATH, 'gprof2dot', 'gprof2dot.py')
+    if imageOutputFile is None:
+        imageOutputFile = os.path.join(paths.SQLMAP_OUTPUT_PATH, "sqlmap_profile.png")
 
-    infoMsg  = "converting profile data to a graph image."
-    logger.info(infoMsg)
+    if os.path.exists(profileOutputFile):
+        os.remove(profileOutputFile)
 
     if os.path.exists(imageOutputFile):
         os.remove(imageOutputFile)
 
-    msg = subprocess.Popen('python %s -f pstats %s | dot -Tpng -o %s' % (graphScriptPath, profileOutputFile, imageOutputFile), shell=True, stderr=subprocess.PIPE).stderr.read()
+    infoMsg  = "profiling the execution into file %s" % profileOutputFile
+    logger.info(infoMsg)
 
-    if msg:
-        errMsg  = "there was an error while converting ('%s'), " % msg.strip()
+    cProfile.run("start()", profileOutputFile)
+
+    infoMsg  = "converting profile data into a graph image, %s" % imageOutputFile
+    logger.info(infoMsg)
+
+    graphScriptPath = os.path.join(paths.SQLMAP_EXTRAS_PATH, 'gprof2dot', 'gprof2dot.py')
+    stderr = subprocess.Popen('python %s -f pstats %s | dot -Tpng -o %s' % (graphScriptPath, profileOutputFile, imageOutputFile), shell=True, stderr=subprocess.PIPE).stderr.read()
+
+    if stderr or not os.path.exists(imageOutputFile):
+        errMsg = "there was an error while converting ('%s')" % stderr.strip()
         errMsg += "but you can still find raw profile data "
         errMsg += "inside file '%s'" % profileOutputFile
         logger.error(errMsg)
-    else:
-        try:
-            if os.name == 'mac':
-                subprocess.call(('open', imageOutputFile))
-            elif os.name == 'posix':
-                subprocess.call(('xdg-open', imageOutputFile))
-            elif os.name == 'nt':
-                subprocess.call(('start', imageOutputFile))
-        except:
-            pass
 
-    if os.path.exists(imageOutputFile):
-        infoMsg  = "done. you can find a graph image inside file '%s'." % imageOutputFile
-        logger.info(infoMsg)
-    else:
-        errMsg  = "there was an error while converting, "
-        errMsg += "but you can still find raw profile data "
-        errMsg += "inside file '%s'" % profileOutputFile
-        logger.error(errMsg)
+        return
+
+    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

lib/core/option.py

@@ -323,7 +323,7 @@ def __setMetasploit():
     if conf.osSmb:
         isAdmin = False
 
-        if "linux" in PLATFORM or "darwin" in PLATFORM:
+        if PLATFORM in ( "posix", "mac" ):
             isAdmin = os.geteuid()
 
             if isinstance(isAdmin, (int, float, long)) and isAdmin == 0:

lib/core/readlineng.py

@@ -61,7 +61,7 @@ if IS_WIN and haveReadline:
 # Thanks to Boyd Waters for this patch.
 uses_libedit = False
 
-if PLATFORM == 'darwin' and haveReadline:
+if PLATFORM == 'mac' and haveReadline:
     import commands
 
     (status, result) = commands.getstatusoutput( "otool -L %s | grep libedit" % _rl.__file__ )

lib/core/settings.py

@@ -23,6 +23,7 @@ Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 """
 
 import logging
+import os
 import subprocess
 import sys
 
@@ -46,7 +47,10 @@ LOGGER.setLevel(logging.WARN)
 
 # System variables
 IS_WIN             = subprocess.mswindows
-PLATFORM           = sys.platform.lower()
+# The name of the operating system dependent module imported. The following
+# names have currently been registered: 'posix', 'nt', 'mac', 'os2', 'ce',
+# 'java', 'riscos'
+PLATFORM           = os.name
 PYVERSION          = sys.version.split()[0]
 
 # Url to update Microsoft SQL Server XML versions file from

lib/takeover/upx.py

@@ -46,15 +46,15 @@ class UPX:
     """
 
     def __initialize(self, srcFile, dstFile=None):
-        if "darwin" in PLATFORM:
+        if PLATFORM == "mac":
             self.__upxPath = "%s/upx/macosx/upx" % paths.SQLMAP_CONTRIB_PATH
 
-        elif "win" in PLATFORM:
+        elif PLATFORM in ( "ce", "nt" ):
             self.__upxTempExe = decloakToMkstemp("%s\upx\windows\upx.exe_" % paths.SQLMAP_CONTRIB_PATH, suffix=".exe")
             self.__upxPath = self.__upxTempExe.name
             self.__upxTempExe.close() #needed for execution rights
 
-        elif "linux" in PLATFORM:
+        elif PLATFORM == "posix":
             self.__upxPath = "%s/upx/linux/upx" % paths.SQLMAP_CONTRIB_PATH
 
         else: