#-*- coding: utf-8 -*-
# (c) 2012 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 hashlib import sha256
from Crypto import Random
from Crypto.Cipher import AES
from Crypto.Util import Counter

# 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 send encrypted data
class SecComSend:
    
    # Save address and password based key, and generate cipher    
    def __init__(self, address, pwd):
        self.addr = address
        self.key = sha256(pwd.encode()).digest()
        self.irv = Random.new().read(AES.block_size)
        ctr = Counter.new(
            AES.block_size*8,
            initial_value=int.from_bytes(self.irv, 'little'))
        self.cipher = AES.new(self.key, AES.MODE_CTR, counter=ctr)
        self._first_time = True

    # Encrypt message and send it
    def send(self, msg):
        if self._first_time:
            self._first_time = False
            tcpsend(self.addr, self.irv)
        tcpsend(self.addr, self.cipher.encrypt(msg))

# Create object to send encrypted data to server, and then send the data
scs = SecComSend(IPaddr(node=host, port=port), passwd)
scs.send("hello".encode())
tcpflush()