#!/usr/bin/env python3
from random import randint
from hashlib import sha256
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad


class dh:
    def __init__(self, p=None, g=2):
        if p is None:
            # NIST's value for p
            self.p = 0xFFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD129024E088A67CC74020BBEA63B139B22514A08798E3404DDEF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7EDEE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3DC2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F83655D23DCA3AD961C62F356208552BB9ED529077096966D670C354E4ABC9804F1746C08CA237327FFFFFFFFFFFFFFFF
        else:
            self.p = p
        self.g = g

        self._private_key = randint(0, self.p)
        self.public_key = pow(self.g, self._private_key, self.p)

    def generate_session_key(self, partners_public_key):
        s = pow(partners_public_key, self._private_key, self.p)
        return sha256(s.to_bytes((s.bit_length() + 7) // 8, "big")).digest()


class client:
    def __init__(self, data_to_send, print_decrypted=False):
        self.data = data_to_send
        self.print_decrypted = print_decrypted

    def start(self):
        self.dh = dh()
        return {"p": self.dh.p, "g": self.dh.g, "public_key": self.dh.public_key}

    def process_packet(self, packet):
        if "public_key" in packet:  # 2nd/Final handshake packet (we sent the first)
            self._session_key = self.dh.generate_session_key(packet["public_key"])[:16]
            cipher = AES.new(self._session_key, AES.MODE_CBC)
            cipher_text = cipher.encrypt(pad(self.data, AES.block_size))
            return {"iv": cipher.iv, "cipher_text": cipher_text}
        else:
            cipher = AES.new(self._session_key, AES.MODE_CBC, packet["iv"])
            plain_text = unpad(cipher.decrypt(packet["cipher_text"]), AES.block_size)
            if self.print_decrypted:
                print(f"Client Decrypted: {plain_text}")
            return None  # Terminate the connection


class echo_server:
    def __init__(self, print_decrypted=False):
        self.echoed = set()
        self.print_decrypted = print_decrypted

    def process_packet(self, packet):
        if "public_key" in packet:
            dh_ = dh(p=packet["p"], g=packet["g"])
            self._session_key = dh_.generate_session_key(packet["public_key"])[:16]
            return {"public_key": dh_.public_key}
        else:
            cipher = AES.new(self._session_key, AES.MODE_CBC, packet["iv"])
            plain_text = unpad(cipher.decrypt(packet["cipher_text"]), AES.block_size)
            self.echoed.add(plain_text)
            if self.print_decrypted:
                print(f"Server Decrypted: {plain_text}")
            cipher = AES.new(self._session_key, AES.MODE_CBC)
            cipher_text = cipher.encrypt(pad(plain_text, AES.block_size))
            return {"iv": cipher.iv, "cipher_text": cipher_text}


class middleman:
    def __init__(self, print_decrypted=False):
        self.found = set()
        self.print_decrypted = print_decrypted

    def process_packet_from_client(self, packet):
        if "public_key" in packet:  # 1st handshake packet
            packet["public_key"] = packet["p"]  # Zero'd session key after modexp

            # Generate the session key ready
            dh_ = dh(p=packet["p"], g=packet["g"])
            self._session_key = dh_.generate_session_key(packet["p"])[:16]

            self.p = packet["p"]  # Save for later when modifying the server's response
            return packet
        else:
            self._decrypt_packet(packet)
            return packet  # No changes

    def process_packet_from_server(self, packet):
        if "public_key" in packet:  # 2nd/Final handshake packet
            packet["public_key"] = self.p  # Zero'd public key after modexp
            return packet
        else:
            self._decrypt_packet(packet)
            return packet  # No changes

    def _decrypt_packet(self, packet):
        cipher = AES.new(self._session_key, AES.MODE_CBC, packet["iv"])
        plain_text = unpad(cipher.decrypt(packet["cipher_text"]), AES.block_size)
        self.found.add(plain_text)
        if self.print_decrypted:
            print(f"Middleman Decrypted: {plain_text}")


class relay:
    def __init__(self, client_side, server_side, mitm=None, print_packets=False):
        self.client = client_side
        self.server = server_side
        self.mitm = mitm
        self.print_packets = print_packets

    # Pass the simulated network packets back and forth between the client & server
    def go(self):
        request = self.client.start()  # Lets get the ball rolling
        while request is not None:
            if self.print_packets:
                print(f"Client Sending: {request}")

            if self.mitm is not None:
                request = self.mitm.process_packet_from_client(request)
                if self.print_packets:
                    print(f"Middleman Sending: {request}")

            response = self.server.process_packet(request)
            if self.print_packets:
                print(f"Server Sending: {response}")

            if self.mitm is not None:
                response = self.mitm.process_packet_from_server(response)
                if self.print_packets:
                    print(f"Middleman Sending: {response}")

            request = self.client.process_packet(response)


def test_normal():
    alices_message = b"Super secret message!"
    alice = client(alices_message)
    bob = echo_server()
    pat = relay(alice, bob)
    pat.go()

    if alices_message not in bob.echoed:
        print("Bob failed to get the message Alice sent!")
        return False
    return True


def test_mitm():
    alices_message = b"Super secret message, Round 2!"
    alice = client(alices_message)
    bob = echo_server()
    eve = middleman()
    pat = relay(alice, bob, mitm=eve)
    pat.go()

    if alices_message not in bob.echoed:
        print("Bob failed to get the message Alice sent!")
        return False
    if alices_message not in eve.found:
        print("Eve wasn't paying attention this morning!")
        return False
    return True


if test_normal():
    print('Alice & Bob can talk "securely"!')
else:
    print("Bob doesn't understand Alice :( (Fail)")

if test_mitm():
    print("Alice & Bob can talk with Eve watching closely!")
else:
    print("Eve is unable to listen in :( (Fail)")
