minor refactoring

stamparm · stamparm · commit bcf9fc6c6f57 · 2012-02-16T09:32:47.000Z
diff --git a/lib/core/common.py b/lib/core/common.py
@@ -61,7 +61,6 @@
 from lib.core.convert import urldecode
 from lib.core.convert import urlencode
 from lib.core.enums import DBMS
-from lib.core.enums import DBMS_DIRECTORY_NAME
 from lib.core.enums import HTTPHEADER
 from lib.core.enums import HTTPMETHOD
 from lib.core.enums import OS
@@ -84,6 +83,7 @@
 from lib.core.settings import INFERENCE_UNKNOWN_CHAR
 from lib.core.settings import UNICODE_ENCODING
 from lib.core.settings import DBMS_DICT
+from lib.core.settings import DBMS_DIRECTORY_DICT
 from lib.core.settings import DESCRIPTION
 from lib.core.settings import DUMMY_SQL_INJECTION_CHARS
 from lib.core.settings import NULL
@@ -1813,9 +1813,7 @@ def getSPLSnippet(dbms, name, **variables):
     Returns content of SPL snippet located inside "procs" directory
     """
 
-    _ = { DBMS.MSSQL: DBMS_DIRECTORY_NAME.MSSQL, DBMS.PGSQL: DBMS_DIRECTORY_NAME.PGSQL }
-
-    filename = os.path.join(paths.SQLMAP_PROCS_PATH, _[dbms], "%s.txt" % name)
+    filename = os.path.join(paths.SQLMAP_PROCS_PATH, DBMS_DIRECTORY_DICT[dbms], "%s.txt" % name)
     checkFile(filename)
     retVal = readCachedFileContent(filename)
 
diff --git a/lib/core/settings.py b/lib/core/settings.py
@@ -14,6 +14,7 @@
 
 from lib.core.enums import CUSTOM_LOGGING
 from lib.core.enums import DBMS
+from lib.core.enums import DBMS_DIRECTORY_NAME
 from lib.core.enums import PLACE
 from lib.core.revision import getRevisionNumber
 
@@ -167,6 +168,8 @@
 SYBASE_ALIASES = ( "sybase", "sybase sql server" )
 DB2_ALIASES = ( "db2", "ibm db2", "ibmdb2" )
 
+DBMS_DIRECTORY_DICT = dict((getattr(DBMS, _), getattr(DBMS_DIRECTORY_NAME, _)) for _ in dir(DBMS) if not _.startswith("_"))
+
 SUPPORTED_DBMS = MSSQL_ALIASES + MYSQL_ALIASES + PGSQL_ALIASES + ORACLE_ALIASES + SQLITE_ALIASES + ACCESS_ALIASES + FIREBIRD_ALIASES + MAXDB_ALIASES + SYBASE_ALIASES + DB2_ALIASES
 SUPPORTED_OS = ( "linux", "windows" )
 
