#!/usr/bin/env python

import string
import urllib
import urllib2
import time
import re
import base64
import sys
from PIL import Image
import pytesseract
import hashlib

## Settings
URL_BASE = "http://51.15.73.163/BOT-WORLD/"
URL_INDEX = URL_BASE + "index.php"
URL_GENERATE = URL_BASE + "generate_captcha.php"

def sha256(text):
    h = hashlib.new('sha256')

    h.update(text.encode("utf-8"))

    return h.hexdigest()

def is_ascii(s):
    return all(ord(c) < 128 for c in s)

def isSolved(title):
    if title == "":
        return False

    if "Incorrect solution" in title:
        return False

    if "Solve captchas to get invitation" in title:
        return False

    return True

def getTitle(data):
    result = re.search(r"<h4>(.+)</h4>", data, re.M|re.I)
    title = result.group(1) if result else ""

    return title

def getWebsiteCode(data):
    data = urllib.unquote(data).decode('utf8')

    result = re.search(r"data=(.+)\'", data, re.M|re.I)
    website_code = result.group(1) if result else ""

    return website_code

def getFlag(data):
    result = re.search(r"(CTF{.+})", data, re.M|re.I)

    flag = result.group(1) if result else ""

    return flag

"""
    Results get/save
"""
def getResults():
    list = {}

    with open('results.txt') as fp:
        for line in fp:
            result = re.search(r"(.+)   (.+)", line, re.M|re.I)

            ## Get website code
            website_code = result.group(1)

            ## Get captcha code
            captcha_code = result.group(2)

            list[website_code] = captcha_code

    return list

def saveResults(website_code, captcha_code):
    f = open("results.txt","a+")

    if f:
        f.write("{} {}\n".format(website_code, captcha_code))

        f.close()

"""

"""
def createSession():
    session = urllib2.build_opener()

    ## Go to index page
    response = session.open(URL_INDEX)
    data = response.read()

    ## Get Cookie
    cookie = response.info().getheader('Set-cookie')
    cookie = cookie[cookie.find("=")+1:cookie.find(";")]

    ## Set Cookie
    session.addheaders.append(("Cookie", "PHPSESSID={}".format(cookie)))

    return session, getWebsiteCode(data)

def downloadImage(session, website_code):
    website_code = urllib.quote(website_code, safe='')

    file_name = "download/{}.png".format(website_code)

    ## Download Picture
    response = session.open(URL_GENERATE + "?data=" + website_code)
    data = response.read()

    ## Save picture
    f = open(file_name, "wb")
    f.write(data)
    f.close()

def removeNoice(website_code):
    website_code = urllib.quote(website_code, safe='')

    infile_path = "download/{}.png".format(website_code)
    outfile_path = "unfilter/{}.png".format(website_code)

    im = Image.open(infile_path)
    im = im.convert('RGB')
    pix = im.load()

    w = im.size[0]
    h = im.size[1]

    gen_im = Image.new("RGB", (w, h), "white")
    gen_pix = gen_im.load()

    # Patch pixels
    for y in range(h):
        for x in range(w):
            if pix[x, y][0] != 15 and pix[x, y][1] != 15 and pix[x, y][2] != 15:
                gen_pix[x, y] = pix[x, y]

    gen_im.save(outfile_path, "PNG")

    return gen_im

"""

"""
def genReverse(results):
    reverse = [{} for x in range(10)]

    ## These cases pytesseract had a hard time to solve, I added this by hand
    reverse[2][43] = "J"
    reverse[3][10] = "l"
    reverse[3][43] = "M"
    reverse[5][43] = "M"

    for website_code in results:
        decode_web = base64.b64decode(website_code)
        captcha_code = results[website_code]

        ## Create results
        reverse[0][ord(decode_web[0])] = captcha_code[0]
        reverse[1][ord(decode_web[1])] = captcha_code[1]
        reverse[2][ord(decode_web[2])] = captcha_code[2]
        reverse[3][ord(decode_web[3])] = captcha_code[3]
        reverse[4][ord(decode_web[4])] = captcha_code[4]
        reverse[5][ord(decode_web[5])] = captcha_code[5]
        reverse[6][ord(decode_web[6])] = captcha_code[6]
        reverse[7][ord(decode_web[7])] = captcha_code[7]
        reverse[8][ord(decode_web[8])] = captcha_code[8]
        reverse[9][ord(decode_web[9])] = captcha_code[9]

    return reverse

def decode(reverse, website_code):
    decode_web = base64.b64decode(website_code)

    ## Debug
    ##print(decode_web)

    captcha_code = reverse[0][ord(decode_web[0])]
    captcha_code += reverse[1][ord(decode_web[1])]
    captcha_code += reverse[2][ord(decode_web[2])]
    captcha_code += reverse[3][ord(decode_web[3])]
    captcha_code += reverse[4][ord(decode_web[4])]
    captcha_code += reverse[5][ord(decode_web[5])]
    captcha_code += reverse[6][ord(decode_web[6])]
    captcha_code += reverse[7][ord(decode_web[7])]
    captcha_code += reverse[8][ord(decode_web[8])]
    captcha_code += reverse[9][ord(decode_web[9])]

    return captcha_code

def test(results, reverse):
    for website_code in results:
        captcha_code = decode(reverse, website_code)

        if captcha_code != results[website_code]:
            print(website_code)
            print(captcha_code)
            print(results[website_code])
            print("-------------")

            print("Oeps something went wrong! :O YOU FAILL!!!!")

            return False

    print("Passed all tests! Good Job! :D")

    return True

def total(reverse):
    c = 0

    for l in reverse:
        c += len(l)

    return c

def analiseResults():
    results = getResults()
    reverse = genReverse(results)

    print(reverse)
    print("total found = {}".format(total(reverse)))

    test(results, reverse)

def runChallange():
    save = False
    solved = 0

    results = getResults()
    reverse = genReverse(results)

    session, website_code = createSession()

    for i in range(100000):
        ## Try to decode with know results
        try:
            ## Generate captcha code
            captcha_code = decode(reverse, website_code)

            save = False
        except:
            ## Not enough results, try to decode the captcha picture

            ## Download image
            downloadImage(session, website_code)

            ## Remove noice out of picture
            im = removeNoice(website_code)

            ## Generate captcha
            captcha_code = pytesseract.image_to_string(im)
            captcha_code = captcha_code if is_ascii(captcha_code) else ""

            save = True

        ## Go to index page
        post = urllib.urlencode({"captcha" : captcha_code, "Submit" : "Submit"})
        response = session.open(URL_INDEX, post)

        data = response.read()
        title = getTitle(data)

        if isSolved(title) == True:
            if save:
                ## Save result
                saveResults(website_code, captcha_code)

                ## Refresh
                results = getResults()
                reverse = genReverse(results)

            solved += 1
        else:
            solved = 0

        print("---------------------------------------")
        print("title = {}".format(title))
        print("website_code = {}".format(website_code))
        print("captcha_code = {}".format(captcha_code))
        print("Learned = {}/620".format(total(reverse)))
        print("Solved = {}".format(solved))
        print("---------------------------------------")

        ## Solved challange! :D
        if "You have proved yourself bot" in title:
            flag = getFlag(title)

            print("Flag = {}".format(flag))
            print("Sha256 = {}".format(sha256(flag)))
            return

        ## Get next website code
        website_code = getWebsiteCode(data)

        ## Debug
        ##print("website_code = {}".format(website_code))

        if website_code == "":
            solved = 0

            ## Create new session
            session, website_code = createSession()

def main(argv):
    if "-a" in argv:
        analiseResults()
    else:
        runChallange()

if __name__ == "__main__":
   main(sys.argv[1:])