#!/usr/bin/env python2.7# @skusec
'''Assumption: randomness of the permutations results in cycles of size < 50(will not always be the case). Start the cycle at the slot we are seeking,so start at the i-th slot in round i, and hope for the best.'''
import websocketimport os, sys, reimport subprocessimport shleximport loggingimport timelogging.basicConfig()
# STATUS CODES (taken from server.py)CHALLENGE = '0'RE_ARRANGED = '1'GIVE_GUESS = '2'GREATE_GUESS = '4'WRONG_GUESS = '5'BYE = '6'FLAG_IS = '7'
def solve_challenge(challenge):    # Native (C) version of the 22-bits MD5 bruteforcer.    cmd = './solve-md5 "%s"' % (challenge)    output = subprocess.check_output(shlex.split(cmd))    return output.strip()
class Game:
    def __init__(self):        self.on_rearranged(0)
    def on_message(self, ws, message):        status, message = message[0], message[1:]        print('>>> [%f] Received message [status=%c]: %s' % (time.time(), status, message))
        if status == CHALLENGE:            solution = solve_challenge(message)            print('>> Solved challenge: %s' % (solution))            self.ws.send(solution)        elif status == RE_ARRANGED:            self.on_rearranged(int(message))        elif status == GIVE_GUESS:            print('>> GIVE GUESS, goal=%d, j=%s' % (self.goal, message))            self.give_guess(int(message))        elif status == GREATE_GUESS:            print('>> **** CORRECT GUESS ****')        elif status == WRONG_GUESS:            actual = int(message)            self.last_actual = actual        elif status == BYE:            print('>> BYE')            exit(1)        elif status == FLAG_IS:            print('>> FLAG: %s' % (message))        else:            print('What??')
        sys.stdout.flush()        sys.stderr.flush()
    def on_rearranged(self, i):        self.goal = i        self.sent_guesses = []        self.last_guess = -1        self.last_actual = -1
    def give_guess(self, j):        if self.last_guess == -1:            self.last_guess = self.goal            self.send_guess(self.goal)        elif self.last_actual != -1:            if self.last_actual in self.sent_guesses:                # We got a cycle, we should have won..                print('WARN: Program logic error, should have gotten slot due to cycle.')                exit(1)            else:                # No cycle yet, keep going..                self.send_guess(self.last_actual)
    def send_guess(self, guess):        self.sent_guesses.append(guess)        self.ws.send(str(guess))
    def start(self):        url = 'ws://213.233.175.130:8998/'        header = ['Cookie: SUCTF_SESSION_ID=86sqbe1nm22j5v3kvnv3ianph6; TEST=207320'] self.ws = websocket.WebSocketApp(url, on_message=self.on_message, header=header)        self.ws.run_forever()
game = Game()game.start()
''' C code for the solver:
#include <stdio.h>#include <string.h>#include <stdint.h>#include <limits.h>#include <stdlib.h>#include <openssl/md5.h>
int main(int argc, char** argv) {
    uint32_t masked_value;    uint32_t guess;    uint32_t upper;    uint32_t current_mask;    uint32_t add_shift;
    MD5_CTX ctx;
    unsigned char md5[16];    char guess_hex[9];
    if (argc != 2) {        return 1;    }
    masked_value = 0;    for (current_mask = 0; current_mask < 22; ++current_mask) {        masked_value |= ((argv[1][current_mask] - 0x30) << (21 - current_mask));    }
    for (guess = 0; guess < UINT_MAX; ++guess) {        memset(md5, 0, 16);        memset(guess_hex, 0, 9);        memset(&ctx, 0, sizeof ctx);
        sprintf(guess_hex, "%08x", guess);        guess_hex[8] = 0;
        MD5_Init(&ctx;;        MD5_Update(&ctx, guess_hex, 8);        MD5_Final(md5, &ctx;;
        upper = 0;        upper |= (md5[0] << 24);        upper |= (md5[1] << 16);        upper |= (md5[2] << 8);        upper |= (md5[3] << 0);
        if ((upper >> 31) == 0) {            continue;        }
        upper = upper >> 10;        if (upper == masked_value) {            printf("%s\n", guess_hex);            return 0;        }    }
    return 1;}
'''
