#!/usr/bin/env python3

import sys

from ast import literal_eval
from manger import Manger
from socket import socket, AF_INET, SOCK_STREAM
from Crypto.Util.number import bytes_to_long, long_to_bytes


def init(io, level):
    # input security level
    io.write(str(level).encode() + b'\n')
    # select RSA
    io.write(b'2\n')
    # show public key
    io.write(b'4\n')
    # show ciphertext
    io.write(b'3\n')
    # select another cipher
    io.write(b'5\n')
    io.flush()
    for i in range(1 + 5 + 7 + 1):
        io.readline()
    n, e = literal_eval(io.readline().decode())
    for i in range(7 + 1):
        io.readline()
    c = bytes_to_long(bytes.fromhex(io.readline().decode().strip()))
    for i in range(7):
        io.readline()
    return n, e, c


def len_oracle(io, c):
    # select RSA
    io.write(b'2\n')
    # set ciphertext
    io.write(b'1\n')
    io.write(long_to_bytes(c).hex().encode() + b'\n')
    # select another cipher
    io.write(b'5\n')
    # select AES
    io.write(b'1\n')
    # show ciphertext
    io.write(b'3\n')
    # select another cipher
    io.write(b'5\n')
    io.flush()
    for i in range(5 + 7 + 1 + 7 + 5 + 7 + 1):
        io.readline()
    ct_len = len(io.readline().strip()) // 2
    for i in range(7):
        io.readline()
    return ct_len


def main(address):
    level = 532
    blocks_count = 8
    while True:
        with socket(AF_INET, SOCK_STREAM) as sock:
            sock.connect(address)
            io = sock.makefile('rwb')
            n, e, c = init(io, level)
            if n.bit_length() != level * 2:
                print('wrong length')
                io.write(b'1\n6\n')
                io.flush()
                continue
            print(n, e, c)
            def oracle(f):
                length = len_oracle(io, f)
                print(length)
                return (length // 16) - 1 >= blocks_count
            manger = Manger(n, e, oracle)
            print(manger.decrypt(c))
            print(manger.oracle_calls)
            break


if __name__ == '__main__':
    ip = sys.argv[1] if len(sys.argv) > 1 else '0.0.0.0'
    main((ip, 44323))
