from telethon._impl.crypto import (
    Side,
    _do_encrypt_data_v2,
    calc_key,
    decrypt_data_v2,
    decrypt_ige,
    encrypt_ige,
    generate_key_data_from_nonce,
)
from telethon._impl.crypto.auth_key import AuthKey


def get_test_auth_key() -> AuthKey:
    return AuthKey.from_bytes(bytes(range(256)))


def get_test_msg_key() -> bytes:
    return bytes(range(16))


def get_test_aes_key_or_iv() -> bytes:
    return bytes(range(32))


def test_calc_client_key() -> None:
    auth_key = get_test_auth_key()
    msg_key = get_test_msg_key()
    expected = (
        b"pN\xd0\x9c\x8bAf\x8a\xe8\xf9\x9d$G8\xf7\x1d\xbd\xdcDF\x9bk\xbdJ\xa8W=\xd0B\xbd\x05\x9e",
        b'M&`\x00\xa5P\xed\xab\xbfL|\xe4\x0f\xd0\x04<\xc9"0\x18L\xd3\x17\xa5\xcc\x9c$\x82\xfd;\x93\x18',
    )
    assert calc_key(auth_key, msg_key, Side.CLIENT) == expected


def test_calc_server_key() -> None:
    auth_key = get_test_auth_key()
    msg_key = get_test_msg_key()
    expected = (
        b"!w%y\x9b$X\x06E\x81t\xa1\xfc\xfb\xc8\x83\x90h\x07\xb1P3\xfd\xd0\xea+Mi\xcf\x9c6N",
        b"f\x9ae8\x91zO\xa5l\xa3#`\xa41\xc9\x16\x0b\xe4\xad\x88q@\x98\r\xab\x91\xce{\xdcG\xff\xbc",
    )
    assert calc_key(auth_key, msg_key, Side.SERVER) == expected


def test_encrypt_client_data_v2() -> None:
    plaintext = b"Hello, world! This data should remain secure!"
    auth_key = get_test_auth_key()
    random_padding = bytes(32)
    expected = b"2\xd1Xn\xa4W\xdf\xc8\xa8\x17)\xd4m\xb5@\x19\xa2\xbf\xd7\xf7D\xf9\xb9lOql\xfd\xc4G}\xb2\xa2\xc1_m\xdb\x85#_\xb9U/\x1d\x84\x07\xc6\xaa\xea\x00\xcc\x84LZ\x1b\xf6\xacD\xb7\x9b^\xdc*#\x86\x8b=`s\xa5\x90\x99,\x0f)u$=V>\xa1\x80\xd2\x18\xeeu|\x9a"

    assert _do_encrypt_data_v2(plaintext, auth_key, random_padding) == expected


def test_decrypt_server_data_v2() -> None:
    ciphertext = b'zq\x83\xc2\xc1\x0eOM\xf9E\xfa\x9a\x9a\xbd5\xe7\xc3\x84\x0ba\xf0E0O9gL\x19\xc0\xe2\txOP\xf6"j\x075)\xd6u\xc9,\xbf\x0b\xfa\x8c\x99\xa7\x9b?9\xc7*]\x9a\x02mC\x1a\xb7@|\xa0N\xccU\x18}lE\xf1xqRN\xdd\x90\xce\xa0.\xd7(\xe1M|\xb1\x8a\xea*caX\xf0\x94Y\xa9Cw\x10\xd8\x94\xc7\x9f6\x8cN\x81d\xb7d~\xa9\x86\x12\xae\xfe\x94,]\x92\x12\x1a\xcb\x8d\xb0-\xcc\xce\xb6m\x0f\x87 \xac\x12\xa0m\xb0X+\xfd\x95[\xe3O6Q\x18\xe3\xba\xb8\xcd\x08\x0c\xe6\xb4[(\xea\xc5m\xcd*)7N'
    auth_key = AuthKey.from_bytes(
        b"].}e\xf4\x9e\xc2\x8b\xd0)\xa8\x87a\xea'\xb8\xa4\xc7\x9f\x12\"e%D>}|Yn\xf3050\xdb!\x07\xe8\x9a\xa9\x97\xc7\xa0\x16J\xb6\x94\x18z\xde\xff\x15k\xd6\xefq\x18\xa1\x96#Gu<\x0e~\x89\xa05K\x8e\xc3d\xf9\x99~q\xbci#\xfb\x86\xe8\xe44\x91\xe0\x10`jl\xe8E\xe2\xfa\x01\x94\tw\xef\n\xa3*\xdfZ\x97\xdb\xf6\xd4(\xec\x044\xd7\x17\xa2\xd3\xad\x19b,\xc0X\x87d!\x13\xc7\x96_\xfb\x86*><\xcb\n\xb9Z\xdd\xdaW\xf8\x92E\xdb\xd7kI#H\xf8\xe9K\xd5\xa7\xc0\xe0\xb8H\x08R<\xfd\x1e\xa8\x0b2\xfe\x9a\xd1\x98\xbc.\x10?\xce\xb7\xd5$\x92\xec\xc0':(gK\xc9#\xee\xe5\x92e\xab\x17\xa0\x02\xdf\x1fJ\xa2\xc5\x9b\x81\x9a^^\x1d\x10^\xc1\x173o\\v\xc6\xb1\x87\x03}KBp\xce\xe9\xcc!\x07\x1d\x97\xe9\xbc\xa2 \xc6\xd7\xb0\x1b\x99\x8c\xf2\xe5\xcd\xb9\xa5\x0e\xcd\xa1\x85*6\xe65i\x0c\x8e"
    )
    expected = b"\xfc\x82j\x02$\x8b(\xfd`\xf2\xc4\x82$C\xadh\x01\xf0\xc1\xc2\x91\x8b0^\x02\x00\x00\x00X\x00\x00\x00\xdc\xf8\xf1s\x02\x00\x00\x00\x01\xa8\xc1\xc2\x91\x8b0^\x01\x00\x00\x00\x1c\x00\x00\x00\x08\t\xc2\x9e\xc4\xfd3\xad\x91\x8b0^\x18\xa8\x8e\xa6\x07\xeeX\x16\xfc\x82j\x02$\x8b(\xfd\x01\xcc\xc1\xc2\x91\x8b0^\x02\x00\x00\x00\x14\x00\x00\x00\xc5sw4\xc4\xfd3\xad\x91\x8b0^d\x080\x00\x00\x00\x00\x00\xfc\xe6g\x04\xa3\xcd\x8e\xe9\xd0\xaeo\xabg,`\xc0J?\x1f\xd4I\x0eQ\xf6"

    assert decrypt_data_v2(ciphertext, auth_key) == expected


def test_key_from_nonce() -> None:
    server_nonce = int.from_bytes(bytes(range(16)))
    new_nonce = int.from_bytes(bytes(range(32)))

    (key, iv) = generate_key_data_from_nonce(server_nonce, new_nonce)
    assert (
        key
        == b'\x07X\xf1S;a]$\xf6\xe8\xa9Jo\xcb\xee\nU\xea\xab"\x17\xd7)\\\xa9!=\x1a-}\x16\xa6'
    )
    assert (
        iv
        == b"Z\x84\x10\x8e\x98\x05el\xe8d\x07\x0e\x16nb\x18\xf6x>\x85\x11G\x1aZ\xb7\x80,\xf2\x00\x01\x02\x03"
    )


def test_verify_ige_encryption() -> None:
    plaintext = get_test_aes_key_or_iv()
    key = get_test_aes_key_or_iv()
    iv = get_test_aes_key_or_iv()
    expected = b"\xe2\x81\x12\xa5>\\\x89\xc7\xb1\xea\x80q\xc13i\x9f\xd4\xe8k&\xc4\xba\xc9\xfcZ\xf1\xab\x8c\xe2zD\xa4"
    assert encrypt_ige(plaintext, key, iv) == expected


def test_verify_ige_decryption() -> None:
    ciphertext = get_test_aes_key_or_iv()
    key = get_test_aes_key_or_iv()
    iv = get_test_aes_key_or_iv()
    expected = b"\xe5wz\xfa\xcd{,\x16\xf7\xac@\xca\xe6\x1e\xf6\x03\xfe\xe6\t\x8f\xb8\xa8\x86\n\xb9\xeeg,\xd7\xe5\xba\xcc"
    assert decrypt_ige(ciphertext, key, iv) == expected