From b3a094b9d67ec4355ab659c1891e58c9ff139780 Mon Sep 17 00:00:00 2001
From: Miroslav Stampar <miroslav.stampar@gmail.com>
Date: Fri, 3 Dec 2010 22:44:29 +0000
Subject: [PATCH] fix for a bug reported by ToR (when resuming:
 queries[kb.dbms] -> KeyError: u'mysql')

---
 lib/core/common.py  | 25 ++++++++++---------------
 lib/core/data.py    | 16 ++++++++++++++++
 lib/core/session.py |  5 +++--
 3 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/lib/core/common.py b/lib/core/common.py
index 468f7cc11..ed9ec9df5 100644
--- a/lib/core/common.py
+++ b/lib/core/common.py
@@ -38,6 +38,7 @@ from xml.sax import parse
 from extra.cloak.cloak import decloak
 from lib.contrib import magic
 from lib.core.data import conf
+from lib.core.data import dbmsDict
 from lib.core.data import kb
 from lib.core.data import logger
 from lib.core.data import paths
@@ -59,13 +60,6 @@ from lib.core.settings import SITE
 from lib.core.settings import SQL_STATEMENTS
 from lib.core.settings import SUPPORTED_DBMS
 from lib.core.settings import VERSION_STRING
-from lib.core.settings import MSSQL_ALIASES
-from lib.core.settings import MYSQL_ALIASES
-from lib.core.settings import PGSQL_ALIASES
-from lib.core.settings import ORACLE_ALIASES
-from lib.core.settings import SQLITE_ALIASES
-from lib.core.settings import ACCESS_ALIASES
-from lib.core.settings import FIREBIRD_ALIASES
 from lib.core.settings import DUMP_NEWLINE_MARKER
 from lib.core.settings import DUMP_CR_MARKER
 from lib.core.settings import DUMP_DEL_MARKER
@@ -733,14 +727,6 @@ def parseTargetDirect():
         errMsg += "or 'access://DATABASE_FILEPATH'"
         raise sqlmapSyntaxException, errMsg
 
-    dbmsDict = { DBMS.MSSQL: [MSSQL_ALIASES, "python-pymssql", "http://pymssql.sourceforge.net/"],
-                 DBMS.MYSQL: [MYSQL_ALIASES, "python-mysqldb", "http://mysql-python.sourceforge.net/"],
-                 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-pysqlite2", "http://pysqlite.googlecode.com/"],
-                 DBMS.ACCESS: [ACCESS_ALIASES, "python-pyodbc", "http://pyodbc.googlecode.com/"],
-                 DBMS.FIREBIRD: [FIREBIRD_ALIASES, "python-kinterbasdb", "http://kinterbasdb.sourceforge.net/"] }
-
     for dbmsName, data in dbmsDict.items():
         if conf.dbms in data[0]:
             try:
@@ -1642,3 +1628,12 @@ def trimAlphaNum(value):
 
 def isNumPosStrValue(value):
     return value and isinstance(value, basestring) and value.isdigit() and value != "0"
+
+def aliasToDbmsEnum(value):
+    retVal = None
+    for key, item in dbmsDict.items():
+        if value in item[0]:
+            retVal = key
+            break
+    return retVal
+
diff --git a/lib/core/data.py b/lib/core/data.py
index 6d85c195f..216e491ba 100644
--- a/lib/core/data.py
+++ b/lib/core/data.py
@@ -8,7 +8,15 @@ See the file 'doc/COPYING' for copying permission
 """
 
 from lib.core.datatype import advancedDict
+from lib.core.enums import DBMS
 from lib.core.settings import LOGGER
+from lib.core.settings import MSSQL_ALIASES
+from lib.core.settings import MYSQL_ALIASES
+from lib.core.settings import PGSQL_ALIASES
+from lib.core.settings import ORACLE_ALIASES
+from lib.core.settings import SQLITE_ALIASES
+from lib.core.settings import ACCESS_ALIASES
+from lib.core.settings import FIREBIRD_ALIASES
 
 # sqlmap paths
 paths = advancedDict()
@@ -25,3 +33,11 @@ queries = {}
 
 # logger
 logger = LOGGER
+
+dbmsDict = { DBMS.MSSQL: [MSSQL_ALIASES, "python-pymssql", "http://pymssql.sourceforge.net/"],
+             DBMS.MYSQL: [MYSQL_ALIASES, "python-mysqldb", "http://mysql-python.sourceforge.net/"],
+             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-pysqlite2", "http://pysqlite.googlecode.com/"],
+             DBMS.ACCESS: [ACCESS_ALIASES, "python-pyodbc", "http://pyodbc.googlecode.com/"],
+             DBMS.FIREBIRD: [FIREBIRD_ALIASES, "python-kinterbasdb", "http://kinterbasdb.sourceforge.net/"] }
diff --git a/lib/core/session.py b/lib/core/session.py
index d573624b9..11338d549 100644
--- a/lib/core/session.py
+++ b/lib/core/session.py
@@ -9,6 +9,7 @@ See the file 'doc/COPYING' for copying permission
 
 import re
 
+from lib.core.common import aliasToDbmsEnum
 from lib.core.common import dataToSessionFile
 from lib.core.common import formatFingerprintString
 from lib.core.common import readInput
@@ -420,10 +421,10 @@ def resumeConfKb(expression, url, value):
             test = readInput(message, default="N")
 
             if not test or test[0] in ("n", "N"):
-                kb.dbms = dbms
+                kb.dbms = aliasToDbmsEnum(dbms)
                 kb.dbmsVersion = dbmsVersion
         else:
-            kb.dbms = dbms
+            kb.dbms = aliasToDbmsEnum(dbms)
             kb.dbmsVersion = dbmsVersion
 
     elif expression == "OS" and url == conf.url: