#!/usr/bin/env python3

import socket
from Crypto.Cipher import AES
from binascii import hexlify, unhexlify
import random

from Crypto.PublicKey import ECC
HOST = 'localhost' 
PORT1 = 20206
PORT2 = 20306



def debug_secure_channel(s1, s2, data):
    data = data.rstrip('\n')   # remove trailing newline
    if len(data) >= 1024:
        print("from {} to {}: `{}...'".format(s1, s2, data[:1024]))
    else:
        print("from {} to {}: `{}'".format(s1, s2, data))

    iv, ciphertext, mac = data.split(";")
    assert (len(iv) == 16 * 2)                # a hexlified byte is two bytes long, the IV should be 16 bytes
    assert (len(ciphertext) % (16 * 2) == 0)  # a hexlified byte is two bytes long, AES block size is 128 bit (16 byte)
    assert (len(mac) == 16 * 2)               # a quite short MAC. Hint: you still don't want to brute force it!
    print("relayed {} encrypted blocks of payload".format(len(unhexlify(ciphertext)) // AES.block_size))


def main():
    s1 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s1.connect((HOST, PORT1))
    s1f = s1.makefile("rw")  # file abstraction for the sockets

    s2 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s2.connect((HOST, PORT2))
    s2f = s2.makefile("rw")

    # First step of Diffie Hellman (compare slides)
    data = s1f.readline()
    print("from s1 to s2: '{}'".format(data))  # you should see the newline at the end printed

    p, g, X = data.split(",")
    p = int(p)
    g = int(g)
    X = int(X)
    print("p: {}  g: {}  X: {}".format(p, g, X))
    a = random.randint(1, p-1)
    eve_X = pow(g, a, mod=p)
    eve_data = f"{p},{g},{eve_X}\n"
    # # forward to Bob
    s2f.write(eve_data)
    s2f.flush()

    # # # Receive answer (second step of DH)
    data = s2f.readline()
    print("from s2 to s1: '{}'".format(data))

    # # # Forward received
    # s1f.write(data)  # the data should still have a b'\n' at the end
    # s1f.flush()

    # # # Alice sends something to Bob
    # data = s1f.readline()
    # print("S1f {}".format(data))
    # debug_secure_channel("Alice", "Bob", data)

    # # # Forward received data to Bob
    # s2f.write(data)
    # s2f.flush()
    # data = s2f.readline()
    # debug_secure_channel("Bob", "Alice", data)
    # print(f"get data from bob: {data}")
    # s1f.write(data)
    # s1f.flush()
    # data = s1f.readline()
    # debug_secure_channel("Bob", "Alice", data)
    # print(f"get data from alice: {data}")
    # # "alice sending hey bob, plz send me my f14g"
    # s2f.write(data)
    # s2f.flush()
    # data = s1f.readline()
    # print(f"get data from alice: {data}")
    # s2f.write(data)
    # s2f.flush()
    # data = s2f.readline()
    # print(f"get data from bob :{data}")
    # s1f.write(data)
    # s1f.flush()
    # data = s1f.readline()
    # print(f"get data from alice again: {data}")
    s1f.close()
    s2f.close()
    s1.close()
    s2.close()


if __name__ == "__main__":
    main()