Cleaner version checking

stamparm · stamparm · commit 060eac110a70 · 2013-02-06T10:28:17.000+01:00
diff --git a/_sqlmap.py b/_sqlmap.py
diff --git a/lib/utils/versioncheck.py b/lib/utils/versioncheck.py
@@ -0,0 +1,13 @@
+#!/usr/bin/env python
+
+"""
+Copyright (c) 2006-2013 sqlmap developers (http://sqlmap.org/)
+See the file 'doc/COPYING' for copying permission
+"""
+
+import sys
+
+PYVERSION = sys.version.split()[0]
+
+if PYVERSION >= "3" or PYVERSION < "2.6":
+    exit("[CRITICAL] incompatible Python version detected ('%s'). For successfully running sqlmap you'll have to use version 2.6 or 2.7 (visit 'http://www.python.org/download/')" % PYVERSION)
diff --git a/sqlmap.py b/sqlmap.py
@@ -5,13 +5,138 @@
 See the file 'doc/COPYING' for copying permission
 """
 
+import bdb
+import os
 import sys
+import time
+import traceback
+import warnings
 
-PYVERSION = sys.version.split()[0]
+warnings.filterwarnings(action="ignore", message=".*was already imported", category=UserWarning)
+warnings.filterwarnings(action="ignore", category=DeprecationWarning)
 
-if PYVERSION >= "3" or PYVERSION < "2.6":
-    exit("[CRITICAL] incompatible Python version detected ('%s'). For successfully running sqlmap you'll have to use version 2.6 or 2.7 (visit 'http://www.python.org/download/')" % PYVERSION)
-elif __name__ == "__main__":
-    from _sqlmap import main
-    from lib.controller.controller import start  # needed for proper working of --profile switch
+from lib.utils import versioncheck
+from lib.controller.controller import start
+from lib.core.common import banner
+from lib.core.common import dataToStdout
+from lib.core.common import getUnicode
+from lib.core.common import setPaths
+from lib.core.common import weAreFrozen
+from lib.core.data import cmdLineOptions
+from lib.core.data import conf
+from lib.core.data import kb
+from lib.core.data import logger
+from lib.core.data import paths
+from lib.core.common import unhandledExceptionMessage
+from lib.core.exception import SqlmapBaseException
+from lib.core.exception import SqlmapSilentQuitException
+from lib.core.exception import SqlmapUserQuitException
+from lib.core.option import initOptions
+from lib.core.option import init
+from lib.core.profiling import profile
+from lib.core.settings import LEGAL_DISCLAIMER
+from lib.core.testing import smokeTest
+from lib.core.testing import liveTest
+from lib.parse.cmdline import cmdLineParser
+from lib.utils.api import setRestAPILog
+from lib.utils.api import StdDbOut
+
+def modulePath():
+    """
+    This will get us the program's directory, even if we are frozen
+    using py2exe
+    """
+
+    return os.path.dirname(getUnicode(sys.executable if weAreFrozen() else __file__, sys.getfilesystemencoding()))
+
+def main():
+    """
+    Main function of sqlmap when running from command line.
+    """
+
+    try:
+        paths.SQLMAP_ROOT_PATH = modulePath()
+        setPaths()
+
+        # Store original command line options for possible later restoration
+        cmdLineOptions.update(cmdLineParser().__dict__)
+        initOptions(cmdLineOptions)
+
+        if hasattr(conf, "api"):
+            # Overwrite system standard output and standard error to write
+            # to an IPC database
+            sys.stdout = StdDbOut(conf.taskid, messagetype="stdout")
+            sys.stderr = StdDbOut(conf.taskid, messagetype="stderr")
+            setRestAPILog()
+
+        banner()
+
+        dataToStdout("[!] legal disclaimer: %s\n\n" % LEGAL_DISCLAIMER, forceOutput=True)
+        dataToStdout("[*] starting at %s\n\n" % time.strftime("%X"), forceOutput=True)
+
+        init()
+
+        if conf.profile:
+            profile()
+        elif conf.smokeTest:
+            smokeTest()
+        elif conf.liveTest:
+            liveTest()
+        else:
+            start()
+
+    except SqlmapUserQuitException:
+        errMsg = "user quit"
+        logger.error(errMsg)
+
+    except (SqlmapSilentQuitException, bdb.BdbQuit):
+        pass
+
+    except SqlmapBaseException as e:
+        e = getUnicode(e)
+        logger.critical(e)
+        sys.exit(1)
+
+    except KeyboardInterrupt:
+        print
+        errMsg = "user aborted"
+        logger.error(errMsg)
+
+    except EOFError:
+        print
+        errMsg = "exit"
+        logger.error(errMsg)
+
+    except SystemExit:
+        pass
+
+    except:
+        print
+        errMsg = unhandledExceptionMessage()
+        logger.critical(errMsg)
+        traceback.print_exc()
+
+    finally:
+        dataToStdout("\n[*] shutting down at %s\n\n" % time.strftime("%X"), forceOutput=True)
+
+        kb.threadContinue = False
+        kb.threadException = True
+
+        if conf.get("hashDB"):
+            try:
+                conf.hashDB.flush(True)
+            except KeyboardInterrupt:
+                pass
+
+        if hasattr(conf, "api"):
+            try:
+                conf.database_cursor.disconnect()
+            except KeyboardInterrupt:
+                pass
+
+        # Reference: http://stackoverflow.com/questions/1635080/terminate-a-multi-thread-python-program
+        if conf.get("threads", 0) > 1 or conf.get("dnsServer"):
+            os._exit(0)
+
+if __name__ == "__main__":
     main()
