Some code refactoring

lib/takeover/abstraction.py

@@ -75,17 +75,17 @@ class Abstraction(Web, UDF, XP_cmdshell):
         return safechardecode(retVal)
 
     def runCmd(self, cmd):
-        getOutput = None
+        choice = None
 
         if not self.alwaysRetrieveCmdOutput:
             message = "do you want to retrieve the command standard "
             message += "output? [Y/n/a] "
-            getOutput = readInput(message, default="Y")
+            choice = readInput(message, default='Y')
 
-            if getOutput in ("a", "A"):
+            if choice in ('a', 'A'):
                 self.alwaysRetrieveCmdOutput = True
 
-        if not getOutput or getOutput in ("y", "Y") or self.alwaysRetrieveCmdOutput:
+        if not choice or choice in ('y', 'Y') or self.alwaysRetrieveCmdOutput:
             output = self.evalCmd(cmd)
 
             if output:
@@ -166,9 +166,8 @@ class Abstraction(Web, UDF, XP_cmdshell):
             msg += "statements as another DBMS user since you provided the "
             msg += "option '--dbms-creds'. If you are DBA, you can enable it. "
             msg += "Do you want to enable it? [Y/n] "
-            choice = readInput(msg, default="Y")
 
-            if not choice or choice in ("y", "Y"):
+            if readInput(msg, default='Y', boolean=True):
                 expression = getSQLSnippet(DBMS.MSSQL, "configure_openrowset", ENABLE="1")
                 inject.goStacked(expression)
 

lib/takeover/udf.py

@@ -42,12 +42,8 @@ class UDF:
     def _askOverwriteUdf(self, udf):
         message = "UDF '%s' already exists, do you " % udf
         message += "want to overwrite it? [y/N] "
-        output = readInput(message, default="N")
 
-        if output and output[0] in ("y", "Y"):
-            return True
-        else:
-            return False
+        return readInput(message, default='N', boolean=True)
 
     def _checkExistUdf(self, udf):
         logger.info("checking if UDF '%s' already exist" % udf)
@@ -327,12 +323,12 @@ class UDF:
 
         msg = "do you want to call your injected user-defined "
         msg += "functions now? [Y/n/q] "
-        choice = readInput(msg, default="Y")
+        choice = readInput(msg, default='Y').strip().upper()
 
-        if choice[0] in ("n", "N"):
+        if choice == 'N':
             self.cleanup(udfDict=self.udfs)
             return
-        elif choice[0] in ("q", "Q"):
+        elif choice == 'Q':
             self.cleanup(udfDict=self.udfs)
             raise SqlmapUserQuitException
 
@@ -347,9 +343,9 @@ class UDF:
             msg += "\n[q] Quit"
 
             while True:
-                choice = readInput(msg)
+                choice = readInput(msg).strip().upper()
 
-                if choice and choice[0] in ("q", "Q"):
+                if choice == 'Q':
                     break
                 elif isinstance(choice, basestring) and choice.isdigit() and int(choice) > 0 and int(choice) <= len(udfList):
                     choice = int(choice)
@@ -390,9 +386,8 @@ class UDF:
             cmd = cmd[:-1]
             msg = "do you want to retrieve the return value of the "
             msg += "UDF? [Y/n] "
-            choice = readInput(msg, default="Y")
 
-            if choice[0] in ("y", "Y"):
+            if readInput(msg, default='Y', boolean=True):
                 output = self.udfEvalCmd(cmd, udfName=udfToCall)
 
                 if output:
@@ -403,9 +398,8 @@ class UDF:
                 self.udfExecCmd(cmd, udfName=udfToCall, silent=True)
 
             msg = "do you want to call this or another injected UDF? [Y/n] "
-            choice = readInput(msg, default="Y")
 
-            if choice[0] not in ("y", "Y"):
+            if not readInput(msg, default='Y', boolean=True):
                 break
 
         self.cleanup(udfDict=self.udfs)

lib/takeover/web.py

@@ -202,9 +202,8 @@ class Web:
         if not kb.absFilePaths:
             message = "do you want sqlmap to further try to "
             message += "provoke the full path disclosure? [Y/n] "
-            getOutput = readInput(message, default="Y")
 
-            if getOutput in ("y", "Y"):
+            if readInput(message, default='Y', boolean=True):
                 headers = {}
                 been = set([conf.url])
 
@@ -391,9 +390,8 @@ class Web:
 
                     message = "do you want to try the same method used "
                     message += "for the file stager? [Y/n] "
-                    getOutput = readInput(message, default="Y")
 
-                    if getOutput in ("y", "Y"):
+                    if readInput(message, default='Y', boolean=True):
                         self._webFileInject(backdoorContent, backdoorName, directory)
                     else:
                         continue

lib/takeover/xp_cmdshell.py

@@ -255,9 +255,8 @@ class XP_cmdshell:
                 message = "xp_cmdshell extended procedure does not seem to "
                 message += "be available. Do you want sqlmap to try to "
                 message += "re-enable it? [Y/n] "
-                choice = readInput(message, default="Y")
 
-                if not choice or choice in ("y", "Y"):
+                if readInput(message, default='Y', boolean=True):
                     self._xpCmdshellConfigure(1)
 
                     if self._xpCmdshellCheck():