minor update and refactoring

lib/core/common.py

@@ -61,6 +61,7 @@ from lib.core.convert import unicodeencode
 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
@@ -1807,12 +1808,14 @@ def parseXmlFile(xmlFile, handler):
     parse(stream, handler)
     stream.close()
 
-def getSPLSnippet(name, **variables):
+def getSPLSnippet(dbms, name, **variables):
     """
     Returns content of SPL snippet located inside "procs" directory
     """
 
-    filename = os.path.join(paths.SQLMAP_PROCS_PATH, "%s.txt" % name)
+    _ = { DBMS.MSSQL: DBMS_DIRECTORY_NAME.MSSQL, DBMS.PGSQL: DBMS_DIRECTORY_NAME.PGSQL }
+
+    filename = os.path.join(paths.SQLMAP_PROCS_PATH, _[dbms], "%s.txt" % name)
     checkFile(filename)
     retVal = readCachedFileContent(filename)
 

lib/core/enums.py

@@ -26,6 +26,7 @@ class SORT_ORDER:
 
 class DBMS:
     ACCESS = "Microsoft Access"
+    DB2 = "IBM DB2"    
     FIREBIRD = "Firebird"
     MAXDB = "SAP MaxDB"
     MSSQL = "Microsoft SQL Server"
@@ -34,7 +35,18 @@ class DBMS:
     PGSQL = "PostgreSQL"
     SQLITE = "SQLite"
     SYBASE = "Sybase"
-    DB2 = "IBM DB2"    
+
+class DBMS_DIRECTORY_NAME:
+    ACCESS = "access"
+    DB2 = "db2"
+    FIREBIRD = "firebird"
+    MAXDB = "maxdb"
+    MSSQL = "mssqlserver"
+    MYSQL = "mysql"
+    ORACLE = "oracle"
+    PGSQL = "postgresql"
+    SQLITE = "sqlite"
+    SYBASE = "sybase"
 
 class CUSTOM_LOGGING:
     PAYLOAD = 9