#!/usr/bin/env python3
from Crypto.Cipher import AES

def dumpNum(n):
    s = ""
    for v in n:
        if len(s) > 0:
            s = s + " "
        s = s + str(v)
    print(s)

def xor_bytes(a,b):
    z = []
    for v in a:
        z.append(v ^ b[len(z)])
    return bytes(z)

def uInt16ToBytes(a):
    return bytes([a >> 8, a & 0xff])

def getNonce(msgid, sessNonce, seccnt):
    return bytes([msgid]) + sessNonce + bytes([0,0]) + uInt16ToBytes(seccnt)

def crypt_data(data, msgid, sessNonce, seccnt, key):
    cipher = AES.new(key, AES.MODE_ECB)
    nonce = getNonce(msgid, sessNonce, seccnt)
    blkid = 1
    o = cipher.encrypt(bytes([1]) + nonce + uInt16ToBytes(blkid))
    return xor_bytes(data, o)

def compute_auth(data, msgid, sessNonce, seccnt, key):
    cipher = AES.new(key, AES.MODE_ECB)
    nonce = getNonce(msgid, sessNonce, seccnt)
    padded_data_length = 16
    padded_data = data.ljust(padded_data_length, b'\0')
    o = cipher.encrypt(bytes([9]) + nonce + uInt16ToBytes(len(data)))
    #TODO: Implement for loop for chunked data
    o = cipher.encrypt(xor_bytes(o, padded_data))
    xv = cipher.encrypt(bytes([1]) + nonce + bytes([0,0]))
    return xor_bytes(o[0:4], xv)

def sendMsg(msgId, data):
    global remote_nonce, key
    getMsg(msgId, data, remote_nonce, key)

def getMsg(msgId, data, remote_nonce, key):
    seccnt = 1
    cd = crypt_data(data, msgId, remote_nonce, seccnt, key)
    bseccnt = uInt16ToBytes(seccnt)
    authval = compute_auth(data, msgId, remote_nonce, seccnt, key)
    seccnt += 1
#    dumpNum(bytes([msgId]) + cd + bseccnt + authval)
    payload = (bytes([128, msgId]) + cd + bseccnt + authval)
#    print(payload.hex())
    return payload
    

#print("TEST")
#key = bytes.fromhex("3b74f06324ad365bc25916a01c547d1d")
#local_nonce = "12345678"
#remote_nonce = bytes([123,122,83,88,105,62,85,55])
#remote_nonce = bytes.fromhex("d7 37 24 61 da f7 0f df")
#
#padcmd = bytes([1,0,0,0,0,0,0,0])
#dumpNum(crypt_data(padcmd, 130, remote_nonce, 1, key))
#sendMsg(135, padcmd)