# -*- coding: utf-8 -*-

# MIT License
#
# Copyright (c) 2017 Pietro Ferretti, Andrea Jegher
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from pwn import *
import pickle
import IPython

conn = remote('bb8.chal.pwning.xxx', 20811)

# ricevi da alice, manda a bob
# ricevi da bob, manda a alice

measured_from_alice = []

for i in range(599):
    print i, 'a'

    # intercept qubits from Alice
    print conn.recvuntil('(y/N)?')
    conn.sendline('y')
    print conn.recvuntil('(Z/Y)')
    conn.sendline('Z')
    print conn.recvuntil('measured ')
    measured_from_alice.append(int(conn.recvuntil('\n').strip()))

    # send my qubits to Bob
    print conn.recvuntil('(Y)?')
    conn.sendline('Y')
    print conn.recvuntil('(Z/Y)')
    conn.sendline('Z')
    print conn.recvuntil('(-1/1)')
    conn.sendline('1')

    # intercept ACK from Bob
    print conn.recvuntil('(y/N)?')
    conn.sendline('y')
    print conn.recvuntil('(Z/Y)')
    conn.sendline('Z')
    print conn.recvuntil('measured ')
    assert(int(conn.recvuntil('\n').strip()) == 1)

    # send ACK to Alice
    print conn.recvuntil('(Y)?')
    conn.sendline('Y')
    print conn.recvuntil('(Z/Y)')
    conn.sendline('Z')
    print conn.recvuntil('(-1/1)')
    conn.sendline('1')

print 599

# intercept qubits from Alice
print conn.recvuntil('(y/N)?')
conn.sendline('y')
print conn.recvuntil('(Z/Y)')
conn.sendline('Z')
print conn.recvuntil('measured ')
measured_from_alice.append(int(conn.recvuntil('\n').strip()))

# send my qubits to Bob
print conn.recvuntil('(Y)?')
conn.sendline('Y')
print conn.recvuntil('(Z/Y)')
conn.sendline('Z')
print conn.recvuntil('(-1/1)')
conn.sendline('1')

### CONTROLLO BASI

bob_guessed_bases = []
alice_correct_bases = []

for i in range(600):
    print i, 'b'

    # receive guessed bases from Bob
    print conn.recvuntil('(y/N)?')
    conn.sendline('y')
    print conn.recvuntil('(Z/Y)')
    conn.sendline('Z')
    print conn.recvuntil('measured ')
    bob_guessed_bases.append(int(conn.recvuntil('\n').strip()))

    # send my guessed bases (Z) to Alice
    print conn.recvuntil('(Y)?')
    conn.sendline('Y')
    print conn.recvuntil('(Z/Y)')
    conn.sendline('Z')
    print conn.recvuntil('(-1/1)')
    conn.sendline('-1')

    # receive responses from Alice
    print conn.recvuntil('(y/N)?')
    conn.sendline('y')
    print conn.recvuntil('(Z/Y)')
    conn.sendline('Z')
    print conn.recvuntil('measured ')
    alice_correct_bases.append(int(conn.recvuntil('\n').strip()))

    # send our responses to Bob
    print conn.recvuntil('(Y)?')
    conn.sendline('Y')
    print conn.recvuntil('(Z/Y)')
    conn.sendline('Z')
    print conn.recvuntil('(-1/1)')
    if bob_guessed_bases[-1] == -1:
        conn.sendline('1')  # corretto
    else:
        conn.sendline('-1') # sbagliato

# Bob manda le basi guessate 
# le teniamo per la risposta dopo
## misura in base Z
## salviamo in una lista i 1 o -1
## -1 è Z, quindi è la base corretta, e viceversa

# Mandiamo a Alice le nostre basi guessate (sempre Z)
# Alice ci risponde sì/no
# ci salviamo le risposte di alice
## mandiamo sempre Z -1
## riceviamo la risposta
## se 1, è corretto
## se -1, sbagliato
## salviamo i valori corretti in una lista
## correct_values_alice

# mandiamo a bob le risposte sì/no
## se -1, mandiamo Z 1 (corretto)
## altrimenti mandiamo Z -1
## correct_values_bob


### VALORI PER CONTROLLO

# per ogni base corretta, salva il valore del bit
correct_bits_alice = []
for i in range(len(alice_correct_bases)):
    if alice_correct_bases[i] == 1:
        correct_bits_alice.append(measured_from_alice[i])

# per ogni base indovinata da Bob, salva il valore del bit
correct_bits_bob = []
for i in range(len(bob_guessed_bases)):
    if bob_guessed_bases[i] == -1:   # Z è corretto
        correct_bits_bob.append(1)

len_alice = len(correct_bits_alice)
len_bob = len(correct_bits_bob)

print "Bit corretti Alice:", len_alice
print "Bit corretti Bob:", len_bob

for i in range(max(len_alice, len_bob) + 1):
    print i, 'c'

    if i <= (len_bob-1)/2:
        # receive even bits from Bob
        print conn.recvuntil('(y/N)?')
        conn.sendline('y')
        print conn.recvuntil('(Z/Y)')
        conn.sendline('Z')
        print conn.recvuntil('measured ')
        assert(int(conn.recvuntil('\n').strip()) == 1)
    else:
        # receive message from bob
        print conn.recvuntil('message (')
        bob_message = conn.recvuntil(') going')[:-7]
        print bob_message
        print correct_bits_bob
        pickle.dump(bob_message, open('bobmessage', 'w'))
        pickle.dump(correct_bits_bob, open('bobkey', 'w'))

    print conn.recvuntil(')?')
    conn.sendline('Y')

    if i <= (len_alice-1)/2:
        # send even bits to Alice
        print conn.recvuntil('(Z/Y)')
        conn.sendline('Z')
        print conn.recvuntil('(-1/1)')
        conn.sendline(str(correct_bits_alice[i*2]))
    else:
        print conn.recvuntil('(Z/Y)')
        conn.sendline('Z')
        print conn.recvuntil('(-1/1)')
        conn.sendline('1')

    if i <= (len_alice-1)/2:
        # receive ACK from Alice
        print conn.recvuntil('(y/N)?')
        conn.sendline('y')
        print conn.recvuntil('(Z/Y)')
        conn.sendline('Z')
        print conn.recvuntil('measured ')
        assert(int(conn.recvuntil('\n').strip()) == 1)
    else:
        # receive message from alice
        print conn.recvuntil('message (')
        alice_message = conn.recvuntil(') going')[:-7]
        print alice_message
        print correct_bits_alice
        pickle.dump(alice_message, open('alicemessage', 'w'))
        pickle.dump(correct_bits_alice, open('alicekey', 'w'))


    # send ACK to Bob
    print conn.recvuntil(')?')
    conn.sendline('Y')
    print conn.recvuntil('(Z/Y)')
    conn.sendline('Z')
    print conn.recvuntil('(-1/1)')
    conn.sendline('1')

IPython.embed()


# Bob manda i valori concordati agli indici 0 2 4 ...
## assert che i valori sono uguali a correct_values_bob

# Io mando ad Alice i valori concordati agli indici 0, 2, 4 ...
## mando i valori di correct_values_alice

# Alice mi risponde con Z 1
## se è sbagliato EOF (uscirà da solo)

# Mando a Bob Z 1 perché sono tutti corretti

## come sistemare x, y lunghezze diverse


### MESSAGGI CIFRATI AES
## chiave Alice = 128 bit corretti alice indici 1,3,5...
## chiave Bob = stesso

# Alice manda messaggio AES
## intercettiamo messaggio e non lo rimandiamo
## decriptiamo con chiave alice

# Bob manda un messaggio AES
## intercettiamo messaggio e non lo rimandiamo
## decriptiamo con chiave bob

# FINE
