Minor code refactoring

stamparm · stamparm · commit 57eb19377edd · 2014-09-16T09:07:31.000+02:00
diff --git a/lib/core/enums.py b/lib/core/enums.py
@@ -338,3 +338,7 @@ class AUTH_TYPE:
     DIGEST = "digest"
     NTLM = "ntlm"
     PKI = "pki"
+
+class AUTOCOMPLETE_TYPE:
+    SQL = 0
+    OS = 1
diff --git a/lib/core/shell.py b/lib/core/shell.py
@@ -13,6 +13,7 @@
 from lib.core.common import Backend
 from lib.core.data import logger
 from lib.core.data import paths
+from lib.core.enums import AUTOCOMPLETE_TYPE
 from lib.core.enums import OS
 
 def saveHistory():
@@ -47,13 +48,13 @@ def global_matches(self, text):
 
         return matches
 
-def autoCompletion(sqlShell=False, osShell=False):
+def autoCompletion(completion=None):
     # First of all we check if the readline is available, by default
     # it is not in Python default installation on Windows
     if not readline._readline:
         return
 
-    if osShell:
+    if completion == AUTOCOMPLETE_TYPE.OS:
         if Backend.isOs(OS.WINDOWS):
             # Reference: http://en.wikipedia.org/wiki/List_of_DOS_commands
             completer = CompleterNG({
diff --git a/lib/takeover/abstraction.py b/lib/takeover/abstraction.py
@@ -13,6 +13,7 @@
 from lib.core.common import readInput
 from lib.core.data import conf
 from lib.core.data import logger
+from lib.core.enums import AUTOCOMPLETE_TYPE
 from lib.core.enums import DBMS
 from lib.core.exception import SqlmapFilePathException
 from lib.core.exception import SqlmapUnsupportedFeatureException
@@ -116,7 +117,7 @@ def shell(self):
             infoMsg += "'x' or 'q' and press ENTER"
             logger.info(infoMsg)
 
-        autoCompletion(osShell=True)
+        autoCompletion(AUTOCOMPLETE_TYPE.OS)
 
         while True:
             command = None
diff --git a/plugins/generic/custom.py b/plugins/generic/custom.py
@@ -15,6 +15,7 @@
 from lib.core.data import conf
 from lib.core.data import logger
 from lib.core.dicts import SQL_STATEMENTS
+from lib.core.enums import AUTOCOMPLETE_TYPE
 from lib.core.settings import NULL
 from lib.core.settings import PARAMETER_SPLITTING_REGEX
 from lib.core.shell import autoCompletion
@@ -73,7 +74,7 @@ def sqlShell(self):
         infoMsg += "'x' or 'q' and press ENTER"
         logger.info(infoMsg)
 
-        autoCompletion(sqlShell=True)
+        autoCompletion(AUTOCOMPLETE_TYPE.SQL)
 
         while True:
             query = None
