Introducing lib/utils/sqlalchemy.py (Issue #361)

stamparm · stamparm · commit 3e65037a0545 · 2013-04-15T10:33:25.000+02:00
diff --git a/lib/controller/handler.py b/lib/controller/handler.py
@@ -18,6 +18,7 @@
 from lib.core.settings import MAXDB_ALIASES
 from lib.core.settings import SYBASE_ALIASES
 from lib.core.settings import DB2_ALIASES
+from lib.utils.sqlalchemy import SQLAlchemy
 
 from plugins.dbms.mssqlserver import MSSQLServerMap
 from plugins.dbms.mssqlserver.connector import Connector as MSSQLServerConn
diff --git a/lib/core/dicts.py b/lib/core/dicts.py
@@ -127,16 +127,16 @@
 DUMP_REPLACEMENTS = {" ": NULL, "": BLANK}
 
 DBMS_DICT = {
-                DBMS.MSSQL: (MSSQL_ALIASES, "python-pymssql", "http://pymssql.sourceforge.net/"),
-                DBMS.MYSQL: (MYSQL_ALIASES, "python pymysql", "https://github.com/petehunt/PyMySQL/"),
-                DBMS.PGSQL: (PGSQL_ALIASES, "python-psycopg2", "http://initd.org/psycopg/"),
-                DBMS.ORACLE: (ORACLE_ALIASES, "python cx_Oracle", "http://cx-oracle.sourceforge.net/"),
-                DBMS.SQLITE: (SQLITE_ALIASES, "python-sqlite", "http://packages.ubuntu.com/quantal/python-sqlite"),
-                DBMS.ACCESS: (ACCESS_ALIASES, "python-pyodbc", "http://pyodbc.googlecode.com/"),
-                DBMS.FIREBIRD: (FIREBIRD_ALIASES, "python-kinterbasdb", "http://kinterbasdb.sourceforge.net/"),
-                DBMS.MAXDB: (MAXDB_ALIASES, None, None),
-                DBMS.SYBASE: (SYBASE_ALIASES, "python-pymssql", "http://pymssql.sourceforge.net/"),
-                DBMS.DB2: (DB2_ALIASES, "python ibm-db", "http://code.google.com/p/ibm-db/"),
+                DBMS.MSSQL: (MSSQL_ALIASES, "python-pymssql", "http://pymssql.sourceforge.net/", "mssql"),
+                DBMS.MYSQL: (MYSQL_ALIASES, "python pymysql", "https://github.com/petehunt/PyMySQL/", "mysql"),
+                DBMS.PGSQL: (PGSQL_ALIASES, "python-psycopg2", "http://initd.org/psycopg/", "postgresql"),
+                DBMS.ORACLE: (ORACLE_ALIASES, "python cx_Oracle", "http://cx-oracle.sourceforge.net/", "oracle"),
+                DBMS.SQLITE: (SQLITE_ALIASES, "python-sqlite", "http://packages.ubuntu.com/quantal/python-sqlite", "sqlite"),
+                DBMS.ACCESS: (ACCESS_ALIASES, "python-pyodbc", "http://pyodbc.googlecode.com/", "access"),
+                DBMS.FIREBIRD: (FIREBIRD_ALIASES, "python-kinterbasdb", "http://kinterbasdb.sourceforge.net/", "firebird"),
+                DBMS.MAXDB: (MAXDB_ALIASES, None, None, "maxdb"),
+                DBMS.SYBASE: (SYBASE_ALIASES, "python-pymssql", "http://pymssql.sourceforge.net/", "sybase"),
+                DBMS.DB2: (DB2_ALIASES, "python ibm-db", "http://code.google.com/p/ibm-db/", "ibm_db_sa"),
             }
 
 FROM_DUMMY_TABLE = {
diff --git a/lib/core/target.py b/lib/core/target.py
@@ -356,7 +356,7 @@ def _resumeDBMS():
 
     if conf.dbms:
         check = True
-        for aliases, _, _ in DBMS_DICT.values():
+        for aliases, _, _, _ in DBMS_DICT.values():
             if conf.dbms.lower() in aliases and dbms not in aliases:
                 check = False
                 break
diff --git a/lib/utils/sqlalchemy.py b/lib/utils/sqlalchemy.py
@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+
+"""
+Copyright (c) 2006-2013 sqlmap developers (http://sqlmap.org/)
+See the file 'doc/COPYING' for copying permission
+"""
+
+import imp
+import sys
+
+_sqlalchemy = None
+try:
+    f, pathname, desc = imp.find_module("sqlalchemy", sys.path[1:])
+    _sqlalchemy = imp.load_module("sqlalchemy", f, pathname, desc)
+except ImportError:
+    pass
+
+from lib.core.data import conf
+from lib.core.data import logger
+from lib.core.exception import SqlmapConnectionException
+from plugins.generic.connector import Connector as GenericConnector
+
+class SQLAlchemy(GenericConnector):
+    def __init__(self):
+        GenericConnector.__init__(self)
+
+    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()
+
+    def fetchall(self):
+        try:
+            return self.cursor.fetchall()
+        except _sqlalchemy.exc.ProgrammingError, msg:
+            logger.log(logging.WARN if conf.dbmsHandler else logging.DEBUG, "(remote) %s" % msg[1])
+            return None
+
+    def execute(self, query):
+        try:
+            self.cursor = self.connection.execute(query)
+        except (_sqlalchemy.exc.OperationalError, _sqlalchemy.exc.ProgrammingError), msg:
+            logger.log(logging.WARN if conf.dbmsHandler else logging.DEBUG, "(remote) %s" % msg[1])
+        except _sqlalchemy.exc.InternalError, msg:
+            raise SqlmapConnectionException(msg[1])
+
+    def select(self, query):
+        self.execute(query)
+        return self.fetchall()
diff --git a/plugins/generic/connector.py b/plugins/generic/connector.py
@@ -46,8 +46,10 @@ def initCursor(self):
 
     def close(self):
         try:
-            self.cursor.close()
-            self.connector.close()
+            if self.cursor:
+                self.cursor.close()
+            if self.connector:
+                self.connector.close()
         except Exception, msg:
             logger.debug(msg)
         finally:
