from pwn import *
import ast
import random
from math import gcd
from Crypto.Util.number import bytes_to_long, long_to_bytes

from server import PRNG, Card, card_str_to_Card, determine_winner

import tqdm

def sim_round(deck, rng, phi, computer_e, computer_d):
    deck = shuffle(rng, deck)
    print("local hacks 1", rng.seed, rng.index, rng.state)
    while computer_e < 2 or computer_d < 1:
        e_array = []
        for _ in range(6):
            e_array.append(str(rng.getnum()))
        computer_e = int(''.join(e_array))
        if gcd(computer_e, phi) == 1:
            computer_d = pow(computer_e,-1,phi)

    return computer_e, computer_d, deck

def shuffle(rng, deck):
    new_deck = []
    for i in range(len(deck)):
        x = rng.getnum()
        if deck[x] not in new_deck:
            new_deck.append(deck[x])
        elif deck[i] not in new_deck:
            new_deck.append(deck[i])
        else:
            for card in deck:
                if card not in new_deck:
                    new_deck.append(card)
                    break
    return new_deck

def exploit():
    deck = []
    deck_str = []
    for suit in ["Spades", "Hearts", "Diamonds", "Clubs"]:
        for i in range(16):
            deck.append(Card(suit, i))
            deck_str.append(str(deck[-1]))
    deck_str = set(deck_str)

    server_deck = list(deck)

    # r = process(["python3", "server.py"])
    r = remote("crypto.csaw.io", 5001)

    r.recvuntil(b"Here are")
    r.recvuntil(b"--> ")

    p, q = r.recvline().decode().strip().split(", ")
    p = int(p)
    q = int(q)
    phi = (p-1)*(q-1)
    N = p*q

    print("p=", p)
    print("q=", q)

    # player_e = 65537
    # player_d = pow(player_e, -1, phi)
    player_e = 1
    player_d = 1

    print("player_e=", player_e)
    print("player_d=", player_d)

    r.recvuntil(b">>")
    r.sendline(str(player_e).encode())
    r.recvuntil(b">>")
    r.sendline(str(player_d).encode())

    r.recvuntil(b"***")
    roundline = r.recvline()
    print("roundline", roundline)
    print(r.recvuntil(b"--> "))
    deck = r.recvline().strip().decode()
    deck = ast.literal_eval(deck)
    # print("recv deck", deck)

    test_item = deck[0]
    found_seed = None
    for seed in tqdm.trange(2**10):
        prng = PRNG(seed)
        rand_values = []
        for i in range(256):
            rand_values.append(prng.getnum())

        found = False

        for offset in range(249):
            e_array = [str(x) for x in rand_values[offset:offset+6]]
            computer_e = int(''.join(e_array))
            if gcd(computer_e, phi) != 1:
                continue
            computer_d = pow(computer_e, -1, phi)

            test_card = long_to_bytes(pow(test_item, computer_d, N)).decode('iso-8859-1')
            if test_card in deck_str:
                print("found seed", test_card, seed)
                found = True
                found_seed = seed
                break

        if found:
            break

    prng = PRNG(found_seed)
    computer_e, computer_d = -1, 0

    for rndnum in range(10):
        if rndnum > 0:
            print(r.recvuntil(b"***"))
            roundline = r.recvline()
            print("roundline", roundline)
            print(r.recvuntil(b"--> "))
            deck = r.recvline().strip().decode()
            deck = ast.literal_eval(deck)
            # print("recv deck", deck)

        computer_e, computer_d, server_deck = sim_round(server_deck, prng, phi, computer_e, computer_d)
        print(computer_e, computer_d, prng.seed, prng.index, prng.state)

        dec_cards = []
        for enc_card in deck:
            dec_card = long_to_bytes(pow(enc_card, computer_d, N)).decode()
            dec_cards.append(card_str_to_Card(dec_card))

        print([str(c) for c in dec_cards])
        print([str(c) for c in server_deck])

        while True:
            random.shuffle(dec_cards)
            computer_cards, player_cards = dec_cards[0:5], dec_cards[5:10]
            computer_winner, player_winner = determine_winner(computer_cards, player_cards)
            if player_winner and not computer_winner:
                break

        server_deck = list(dec_cards)

        enc_deck = []
        for card in dec_cards:
            enc_deck.append(pow(bytes_to_long(str(card).encode()),computer_e,N))

        for c in enc_deck:
            r.recvuntil(b">> ")
            r.sendline(str(c).encode())

    r.recvuntil(b"HAHAHA")
    r.recvline()

    val = ast.literal_eval(r.recvline().decode())
    print(long_to_bytes(pow(bytes_to_long(val),computer_d,N)))

if __name__ == "__main__":
    exploit()