#!/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
server_ip_address, port = arguments[1], int(arguments[2])

# Server
if server_ip_address == "-l":
    
    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"
    
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.bind((server_ip_address, port))
        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 = calc_hash(K)

        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, AES.MODE_GCM, nonce)
                    padded_text = cipher.decrypt_and_verify(ciphertext, tag)
                    data = unpad(padded_text, 16, style='pkcs7')
                    sys.stdout.buffer.write(data)
                
                except(ValueError, KeyError):
                    sys.stderr.write("Error: integrity check failed.")

# Client
else:
    
    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, port))
        
        # Send A to server
        sock.sendall(A_padded)
        
        # Get B from server
        B = int(sock.recv(384).decode("utf-8"))
        print(B)
        K = pow(B, a, p)
        key = calc_hash(K)
        
        
        pdu_size = 1024
        
        while True:
            pdu = sys.stdin.buffer.read(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)