#-*- coding: utf-8 -*-
# (c) 2015 Anders Andersen
# See http://www.cs.uit.no/~aa/dist/tools/py/COPYING for details

# Load modules
import sys
from noop.ip.tcp import *
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes

# Read arguments (host, port and password)
host = "localhost"
port = 3456
passwd = "pw"
if len(sys.argv) > 1:
    host = sys.argv[1]    
if len(sys.argv) > 2:
    port = int(sys.argv[2])
if len(sys.argv) > 3:
    passwd = sys.argv[3]

# A class to receive encrypted data
class SecComReceive:

    # Save address and password based key
    def __init__(self, address, pwd):
        self.addr = address
        dig = hashes.Hash(hashes.SHA256(), backend=default_backend())
        dig.update(pwd.encode())
        self.key = dig.finalize()
        self._first_time = True

    # Receive cipher text and decrypt it
    def receive(self):
        data = b""
        if self._first_time:
            self._first_time = False
            tmp = tcpreceive(self.addr)
            while len(tmp) < algorithms.AES.block_size//8:
                tmp += tcpreceive(self.addr)
            self.irv = tmp[:algorithms.AES.block_size//8]
            data = tmp[algorithms.AES.block_size//8:]
            cip = Cipher(algorithms.AES(self.key), modes.CTR(self.irv), backend=default_backend())
            self.dec = cip.decryptor()
        if not len(data) > 0:
            data = tcpreceive(self.addr)
        return (self.dec.update(data) + self.dec.finalize())

# Create object to receive data from client, and then receive the data
scr = SecComReceive(IPaddr(node=host, port=port), passwd)
print(scr.receive().decode())