from crawlerhttp import CrawlerHttp, CrawlerHttpSubmission, CrawlerHttpResponse, tokenKeys
import sys
import urllib
import random
import string

def readInput(path):
    inputs = []
    with open(path) as f:
        for line in f:
            inputs = inputs + [line.rstrip()]
    return inputs

redirectCodes = [301, 302, 303, 307, 308]

def attack(submissionList, inputsource, expectedsource, checkFailure = False, checkRedirect = False):
    dataForInput = readInput(inputsource)
    dataExpected = readInput(expectedsource)
    attackFormSubmission = []
    for submission in submissionList:
        print("Attacking " + submission.url)

        for encode in range(0, 2):

            for i in dataForInput:
                duplicate = {}
                duplicate = submission.duplicate()
                duplicate.alterData(i)
                if (encode == 1):
                    # with encoding
                    duplicate.encodeData()
                # TODO - double encoding ??
                #elif (encode == 2):
                #    duplicate.encodeData()
                #    duplicate.encodeData()

                duplicateSubmission = duplicate.submit()
                if (duplicateSubmission is None):
                    print("Attack failure")
                    break
                else:
                    appendToSubmissionResult = False

                    if (checkFailure and int(duplicateSubmission.getCode()) == 500):
                        appendToSubmissionResult = True

                    if (checkRedirect and int(duplicateSubmission.getCode()) in redirectCodes):
                        appendToSubmissionResult = True

                    # check if expected is found. If found, exploit is successful
                    for e in dataExpected:
                        if duplicateSubmission.getPage().find(e) != -1:
                            appendToSubmissionResult = True
                            break

                    if (appendToSubmissionResult):
                        print("Attack Successful")
                        attackFormSubmission = attackFormSubmission + [duplicate]
                        break

            encode = encode + 1

    return attackFormSubmission


def attack_sqli(submissionList, inputsource, expectedsource, checkFailure=False, checkRedirect=False):

    dataForInput = readInput(inputsource)
    dataExpected = readInput(expectedsource)
    attackFormSubmission = []
    for submission in submissionList:
        print("Attacking " + submission.url)

        # initial submission to get the response size
        duplicate = {}
        duplicate = submission.duplicate()
        duplicate.alterData('longtextlongerthanthequerystring')
        duplicateSubmission = duplicate.submit()
        initialPageSize = sys.getsizeof(duplicateSubmission.getPage())

        for encode in range(0, 2):
            for i in dataForInput:
                duplicate = {}
                duplicate = submission.duplicate()
                duplicate.alterData(i)

                if (encode == 1):
                    duplicate.encodeData()
                # TODO - double encoding ??
                # elif (encode == 2):
                #    duplicate.encodeData()
                #    duplicate.encodeData()
                duplicateSubmission = duplicate.submit()
                if (duplicateSubmission is None):
                    print("Attack failure")
                    break
                else:
                    appendToSubmissionResult = False

                    if (checkFailure and int(duplicateSubmission.getCode()) == 500):
                        appendToSubmissionResult = True

                    if (checkRedirect and int(duplicateSubmission.getCode()) in redirectCodes):
                        appendToSubmissionResult = True

                    # check the size of the response. If greater than the initial value, exploit is successful
                    if initialPageSize < sys.getsizeof(duplicateSubmission.getPage()):
                        appendToSubmissionResult = True

                    if (appendToSubmissionResult):
                        print("Attack Successful")
                        attackFormSubmission = attackFormSubmission + [duplicate]
                        break

            encode = encode + 1
    return attackFormSubmission








def find_csrf_token(key):

    # Check for common CSRF token names
    for common_csrf_name in tokenKeys:
        if common_csrf_name.lower() in key.lower():
            print ("Found CSRF token %s" % key)
            return True

    return False

def findCsrf (submission):
    for key, value in submission.data.items():
        # find csrf
        if find_csrf_token(key):
            return value
    return None


def internal_spider(mycrawler, pages):

    crawledPages = []
    for page in pages:
        crawledDocument = mycrawler.send(page.url)
        crawledPages = crawledPages + [crawledDocument]

    return crawledPages

def attack_csrf(mycrawler, pages, submissionList, checkFailure = False, checkRedirect = False):

    attackFormSubmission = []
    for submission in submissionList:
        print("Attacking " + submission.url)

        csrfFound = findCsrf (submission)

        if csrfFound is None:
            continue

        ## here we found the form that has csrf

        # check if the token is fixed, then it is vulnerable
        # refresh the form page and see if the csrf changes

        # 1. get page url
        # 2. load the page again

        pageReloaded = mycrawler.send(submission.url)
        pageReloadedSubmission = pageReloaded.getSubmission()

        for submission2 in pageReloadedSubmission:
            csrfFound2 = findCsrf(submission2)
            if csrfFound2 is None:
                continue
            else:
                break

        if (csrfFound != csrfFound2):   # this means csrf changes, meaning not vulnerable
            continue

        #formURL = submissions[index].duplicate().url  #form action url

        # Test submission of the post form
        # response code 200 means the submission/exploit is successful
        randomString = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(10))

        duplicate = {}
        duplicate = submission.duplicate()
        duplicate.alterData(randomString)  # insert the random string
        duplicateSubmission = duplicate.submit()

        if (duplicateSubmission is None):
            print("Attack failure")
            break
        else:
            appendToSubmissionResult = False

            if int(duplicateSubmission.getCode()) == 200:

                # find any pages that contains the randomString
                # that page will probably the form view page

                # refresh pages
                # temp_pages = internal_spider(mycrawler, pages)
                #
                # for page in temp_pages:
                #     if (randomString in page.resp.text):
                #         formPage = page
                #         appendToSubmissionResult = True
                #         break

                appendToSubmissionResult = True

            if (appendToSubmissionResult):
                print("Attack Successful")
                #attackFormSubmission = attackFormSubmission + [duplicate]

                #attackFormSubmission = attackFormSubmission + [crawledDocument]
                result = []
                result = result + [duplicate] #+ [formPage]
                attackFormSubmission = attackFormSubmission + [result]

    return attackFormSubmission