"""This module holds various utility functions used by the MSL client"""

import base64
import json
import os
import random
import string

from Cryptodome.Cipher import AES
from Cryptodome.Util import Padding


def dumps(data):
    """
    dumps()

    @param data: Data as a dict to use for creating JSON string

    @return: JSON string without whitespace
    """

    return json.dumps(data, separators=(',', ':'))


def msl_encrypt(msl_session, plaintext):
    """
    msl_encrypt()

    @param msl_session: Dict of msl_session created by the client
                        upon initialization
    @param plaintext: Plaintext to encrypt

    @return: JSON byte string of encryption envelope
    """

    cbc_iv = os.urandom(16)
    encryption_envelope = {
        'keyid': '%s_%s' % (
            msl_session['esn'],
            msl_session['session_keys']['sequence_number']
        ),
        'sha256': 'AA==',
        'iv': base64.b64encode(cbc_iv).decode('utf8')
    }

    plaintext = Padding.pad(plaintext.encode('utf8'), 16)
    cipher = AES.new(
        msl_session['session_keys']['encryption_key'],
        AES.MODE_CBC,
        cbc_iv
    )

    ciphertext = cipher.encrypt(plaintext)

    encryption_envelope['ciphertext'] = base64.b64encode(
        ciphertext
    ).decode('utf8')

    return json.dumps(encryption_envelope).encode('utf8')


def generate_esn(prefix):
    """
    generate_esn()

    @param prefix: Prefix of ESN to append generated device ID onto

    @return: ESN to use with MSL API
    """

    return prefix + ''.join(random.choice(
        string.ascii_uppercase + string.digits
    ) for _ in range(30))


def webcrypto_b64decode(b64):
    """
    webcrypto_b64decode()

    @param b64: URL safe encoded base64 string lacking padding
                (most likely from WebCrypto API)

    @return: Bytes from decoding base64 string
    """

    padding = len(b64) % 4
    if padding != 0:
        b64 += '=' * (4 - padding)
    return base64.urlsafe_b64decode(b64.encode('utf8'))