"""AES-256-CBC encryption/decryption and signature utilities for WeCom."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_crypto.ipynb.

# %% auto #0
__all__ = ['pkcs7_pad', 'pkcs7_unpad', 'msg_sig', 'decrypt', 'encrypt']

# %% ../nbs/01_crypto.ipynb #bec3f596
import base64, hashlib, os, struct
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from .config import AES_KEY, CORP_ID

# %% ../nbs/01_crypto.ipynb #82e5b475
def pkcs7_pad(
    data: bytes,  # Data to pad
    bs: int = 32  # Block size; WeCom requires 32
) -> bytes:
    "Pad `data` to a multiple of `bs` bytes using PKCS#7."
    n = bs - len(data) % bs
    return data + bytes([n]) * n

def pkcs7_unpad(
    data: bytes  # PKCS#7-padded bytes
) -> bytes:
    "Remove PKCS#7 padding from `data`."
    return data[: -data[-1]]

# %% ../nbs/01_crypto.ipynb #3251cdcc
def msg_sig(
    *parts: str  # Strings to sort, concatenate, and SHA1-hash
) -> str:
    "WeCom message signature: sort `parts`, join, and return SHA1 hex digest."
    return hashlib.sha1(''.join(sorted(parts)).encode()).hexdigest()

# %% ../nbs/01_crypto.ipynb #038feb5d
def decrypt(
    encrypted: str  # Base64-encoded AES-256-CBC ciphertext from WeCom
) -> tuple:
    "Decrypt a WeCom message; returns `(xml_text, corp_id)`."
    dec = Cipher(algorithms.AES(AES_KEY), modes.CBC(AES_KEY[:16])).decryptor()
    raw = pkcs7_unpad(dec.update(base64.b64decode(encrypted)) + dec.finalize())
    n = struct.unpack('>I', raw[16:20])[0]
    return raw[20:20+n].decode(), raw[20+n:].decode()

def encrypt(
    xml: str  # Plaintext XML to encrypt for a WeCom passive reply
) -> str:
    "Encrypt `xml` using AES-256-CBC for WeCom; returns a Base64 string."
    msg = xml.encode()
    body = pkcs7_pad(os.urandom(16) + struct.pack('>I', len(msg)) + msg + CORP_ID.encode())
    enc = Cipher(algorithms.AES(AES_KEY), modes.CBC(AES_KEY[:16])).encryptor()
    return base64.b64encode(enc.update(body) + enc.finalize()).decode()
