Update for an Issue #361

stamparm · stamparm · commit aed738d6e6d6 · 2013-04-15T14:20:21.000+02:00
diff --git a/lib/controller/handler.py b/lib/controller/handler.py
@@ -8,6 +8,8 @@
 from lib.core.common import Backend
 from lib.core.data import conf
 from lib.core.data import logger
+from lib.core.dicts import DBMS_DICT
+from lib.core.enums import DBMS
 from lib.core.settings import MSSQL_ALIASES
 from lib.core.settings import MYSQL_ALIASES
 from lib.core.settings import ORACLE_ALIASES
@@ -48,16 +50,16 @@ def setHandler():
     """
 
     items = [
-                  ("MySQL", MYSQL_ALIASES, MySQLMap, MySQLConn),
-                  ("Oracle", ORACLE_ALIASES, OracleMap, OracleConn),
-                  ("PostgreSQL", PGSQL_ALIASES, PostgreSQLMap, PostgreSQLConn),
-                  ("Microsoft SQL Server", MSSQL_ALIASES, MSSQLServerMap, MSSQLServerConn),
-                  ("SQLite", SQLITE_ALIASES, SQLiteMap, SQLiteConn),
-                  ("Microsoft Access", ACCESS_ALIASES, AccessMap, AccessConn),
-                  ("Firebird", FIREBIRD_ALIASES, FirebirdMap, FirebirdConn),
-                  ("SAP MaxDB", MAXDB_ALIASES, MaxDBMap, MaxDBConn),
-                  ("Sybase", SYBASE_ALIASES, SybaseMap, SybaseConn),
-                  ("IBM DB2", DB2_ALIASES, DB2Map, DB2Conn),
+                  (DBMS.MYSQL, MYSQL_ALIASES, MySQLMap, MySQLConn),
+                  (DBMS.ORACLE, ORACLE_ALIASES, OracleMap, OracleConn),
+                  (DBMS.PGSQL, PGSQL_ALIASES, PostgreSQLMap, PostgreSQLConn),
+                  (DBMS.MSSQL, MSSQL_ALIASES, MSSQLServerMap, MSSQLServerConn),
+                  (DBMS.SQLITE, SQLITE_ALIASES, SQLiteMap, SQLiteConn),
+                  (DBMS.ACCESS, ACCESS_ALIASES, AccessMap, AccessConn),
+                  (DBMS.FIREBIRD, FIREBIRD_ALIASES, FirebirdMap, FirebirdConn),
+                  (DBMS.MAXDB, MAXDB_ALIASES, MaxDBMap, MaxDBConn),
+                  (DBMS.SYBASE, SYBASE_ALIASES, SybaseMap, SybaseConn),
+                  (DBMS.DB2, DB2_ALIASES, DB2Map, DB2Conn),
             ]
 
     _ = max(_ if (Backend.getIdentifiedDbms() or "").lower() in _[1] else None for _ in items)
@@ -77,7 +79,15 @@ def setHandler():
         if conf.direct:
             logger.debug("forcing timeout to 10 seconds")
             conf.timeout = 10
-            conf.dbmsConnector.connect()
+
+            dialect = DBMS_DICT[name][3]
+            sqlalchemy = SQLAlchemy(dialect=dialect)
+            sqlalchemy.connect()
+
+            if sqlalchemy.connection:
+                conf.dbmsConnector = sqlalchemy
+            else:
+                conf.dbmsConnector.connect()
 
         if handler.checkDbms():
             conf.dbmsHandler = handler
diff --git a/lib/core/common.py b/lib/core/common.py
@@ -137,6 +137,7 @@
 from lib.core.settings import VERSION
 from lib.core.settings import VERSION_STRING
 from lib.core.threads import getCurrentThreadData
+from lib.utils.sqlalchemy import _sqlalchemy
 from thirdparty.clientform.clientform import ParseResponse
 from thirdparty.clientform.clientform import ParseError
 from thirdparty.magic import magic
@@ -1121,10 +1122,15 @@ def parseTargetDirect():
                 elif dbmsName == DBMS.FIREBIRD:
                     import kinterbasdb
             except ImportError:
-                errMsg = "sqlmap requires '%s' third-party library " % data[1]
-                errMsg += "in order to directly connect to the database "
-                errMsg += "%s. Download from '%s'" % (dbmsName, data[2])
-                raise SqlmapMissingDependence(errMsg)
+                if _sqlalchemy and data[3] in _sqlalchemy.dialects.__all__:
+                    pass
+                else:
+                    errMsg = "sqlmap requires '%s' third-party library " % data[1]
+                    errMsg += "in order to directly connect to the database "
+                    errMsg += "%s. You can download it from '%s'" % (dbmsName, data[2])
+                    errMsg += ". Alternative is to use a package 'python-sqlalchemy' "
+                    errMsg += "with support for dialect '%s' installed" % data[3]
+                    raise SqlmapMissingDependence(errMsg)
 
 def parseTargetUrl():
     """
diff --git a/lib/utils/sqlalchemy.py b/lib/utils/sqlalchemy.py
@@ -21,22 +21,27 @@
 from plugins.generic.connector import Connector as GenericConnector
 
 class SQLAlchemy(GenericConnector):
-    def __init__(self):
+    def __init__(self, dialect=None):
         GenericConnector.__init__(self)
 
+        self.dialect = dialect
+
     def connect(self):
-        self.initConnection()
-        try:
-            #_sqlalchemy.dialects.__all__
-            if not self.port and self.db:
-                if "///" not in conf.direct:
-                    conf.direct = conf.direct.replace("//", "///")
-            engine = _sqlalchemy.create_engine(conf.direct, connect_args={'check_same_thread':False})
-            self.connection = engine.connect()
-        except _sqlalchemy.exc.OperationalError, msg:
-            raise SqlmapConnectionException(msg[0])
-
-        self.connected()
+        if _sqlalchemy:
+            self.initConnection()
+
+            try:
+                if not self.port and self.db:
+                    if "///" not in conf.direct:
+                        conf.direct = conf.direct.replace("//", "///", 1)
+                if self.dialect:
+                    conf.direct = conf.direct.replace(conf.dbms, self.dialect)
+                engine = _sqlalchemy.create_engine(conf.direct, connect_args={'check_same_thread':False} if self.dialect == "sqlite" else {})
+                self.connection = engine.connect()
+            except Exception, msg:
+                raise SqlmapConnectionException(msg[0])
+
+            self.connected()
 
     def fetchall(self):
         try:
diff --git a/plugins/generic/connector.py b/plugins/generic/connector.py
@@ -18,6 +18,7 @@ class Connector:
     """
 
     def __init__(self):
+        self.connection = None
         self.connector = None
         self.cursor = None
 
