#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Sep 17 00:00:28 2022

@author: sarthakbabbar
"""
import os
import sys
import socket
import struct

from Crypto.Util.Padding import pad, unpad
from Crypto.Cipher import AES
import hashlib

p = 0x00cc81ea8157352a9e9a318aac4e33ffba80fc8da3373fb44895109e4c3ff6cedcc55c02228fccbd551a504feb4346d2aef47053311ceaba95f6c540b967b9409e9f0502e598cfc71327c5a455e2e807bede1e0b7d23fbea054b951ca964eaecae7ba842ba1fc6818c453bf19eb9c5c86e723e69a210d4b72561cab97b3fb3060b
g = 2

def pad_zero(a):
    a_str = str(a)
    
    if len(a_str) < 384:
        a_str = '0'*(384-len(a_str)) + a_str
    
    return a_str.encode('utf-8')

def calc_hash(key):
    key_string = '%x' % key  
    key_encoded = key_string.encode('utf-8')
    m = hashlib.sha256()
    m.update(key_encoded)
    hashed_string = m.digest()
    return hashed_string[:32]

# Getting command line arguments
arguments = sys.argv
pport, server_ip_address, sport = int(arguments[2]), arguments[3], int(arguments[4])

# Proxy

    
b = -1
while b < 0:
    b = struct.unpack('i', os.urandom(4))[0]
B = pow(g, b, p)
B_padded = pad_zero(B)

server_ip_address = "127.0.0.1"


# get data from Client
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
    sock.bind((server_ip_address, pport))
    sock.listen()
    connection, _ = sock.accept()
    
    #Receive A from client
    A = int(connection.recv(384).decode("utf-8"))

    #Send B to client
    connection.sendall(B_padded)
    
    K = pow(A, b, p)
    key_client = calc_hash(K)
    data = []

    with connection:
        while True:
            
            packed_length = connection.recv(2)
            if len(packed_length)<2:
                break
            unpacked_length = struct.unpack('!H', packed_length)[0]
            nonce = connection.recv(16)
            tag = connection.recv(16)
            
            ciphertext = connection.recv(unpacked_length-32)
            
            if not ciphertext:
                break
            
            try:
                cipher = AES.new(key_client, AES.MODE_GCM, nonce)
                padded_text = cipher.decrypt_and_verify(ciphertext, tag)
                data.append(unpad(padded_text, 16, style='pkcs7'))
                # sys.stdout.buffer.write(data)
                
            except(ValueError, KeyError):
                sys.stderr.write("Error: integrity check failed.")


# Send client data to server
    
a = -1
while a < 0:
    a = struct.unpack('i', os.urandom(4))[0]
A = pow(g, a, p)
A_padded = pad_zero(A)

server_ip_address = str(server_ip_address)

with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
    sock.connect((server_ip_address, sport))
    
    # Send A to server
    sock.sendall(A_padded)
    
    # Get B from server
    B = int(sock.recv(384).decode("utf-8"))
    K = pow(B, a, p)
    key = calc_hash(K)
    pdu_size = 1024
    
    print(data)
    
    for pdu in data:
    # for i in range(0, len(data), pdu_size):
        # pdu = data[i:i+pdu_size]
        
        if len(pdu) == 0:
            break
        
        padded_data = pad(pdu,16,style='pkcs7')
        cipher = AES.new(key, AES.MODE_GCM)
        ciphertext, tag = cipher.encrypt_and_digest(padded_data)
        nonce = cipher.nonce
        formatted_length = struct.pack('!H', 32 + len(ciphertext))
        sock.sendall(formatted_length + nonce + tag + ciphertext)