Repository: Zapix/mtpylon (language: python)

Issue:
Pack/Unpack Unencrypted messages
- [ ] Pack message
- [ ] Unpack message
- [ ] Validate message ID for client
- [ ] Validate message ID for server



Discussion / hints:

Fix patch:
diff --git a/mtpylon/messages/__init__.py b/mtpylon/messages/__init__.py
new file mode 100644
index 0000000..b99370a
--- /dev/null
+++ b/mtpylon/messages/__init__.py
@@ -0,0 +1,16 @@
+# -*- coding: utf-8 -*-
+from .utils import is_encrypted_message, is_unencrypted_message
+from .unencrypted_message import (
+    UnencryptedMessage,
+    unpack_message,
+    pack_message
+)
+
+
+__all__ = [
+    'is_unencrypted_message',
+    'is_encrypted_message',
+    'unpack_message',
+    'pack_message',
+    'UnencryptedMessage',
+]
diff --git a/mtpylon/messages/unencrypted_message.py b/mtpylon/messages/unencrypted_message.py
new file mode 100644
index 0000000..16b044e
--- /dev/null
+++ b/mtpylon/messages/unencrypted_message.py
@@ -0,0 +1,54 @@
+# -*- coding: utf-8 -*-
+import asyncio
+from dataclasses import dataclass
+from typing import Any
+
+from mtpylon.utils import long
+from mtpylon.service_schema import load as load_schema, dump as dump_schema
+from mtpylon.serialization.int import load as load_int, dump as dump_int
+from mtpylon.serialization.long import load as load_long, dump as dump_long
+from mtpylon.serialization import CallableFunc
+
+
+@dataclass
+class UnencryptedMessage:
+    msg_id: long
+    value: Any
+
+
+def unpack(input: bytes) -> UnencryptedMessage:
+    loaded_msg_id = load_long(input[8:])
+    loaded_size = load_int(input[16:])
+    size = loaded_size.value
+    loaded_value = load_schema(input[20:20 + size])
+
+    return UnencryptedMessage(
+        msg_id=loaded_msg_id.value,
+        value=loaded_value.value
+    )
+
+
+def pack(message: UnencryptedMessage):
+    auth_id = dump_long(long(0))
+
+    msg_id = dump_long(message.msg_id)
+
+    value: bytes = b''
+    if isinstance(message.value, CallableFunc):
+        func = message.value.func
+        params = message.value.params
+        value = dump_schema(func, **params)
+    else:
+        value = dump_schema(message.value)
+
+    size = dump_int(len(value))
+
+    return auth_id + msg_id + size + value
+
+
+async def unpack_message(input: bytes) -> UnencryptedMessage:
+    return await asyncio.to_thread(unpack, input)
+
+
+async def pack_message(message: UnencryptedMessage) -> bytes:
+    return await asyncio.to_thread(pack, message)
diff --git a/mtpylon/messages/utils.py b/mtpylon/messages/utils.py
new file mode 100644
index 0000000..e50de93
--- /dev/null
+++ b/mtpylon/messages/utils.py
@@ -0,0 +1,20 @@
+# -*- coding: utf-8 -*-
+from mtpylon.serialization.long import load as load_long
+
+
+def is_unencrypted_message(buffer: bytes) -> bool:
+    """
+    Unencrypted message starts with auth_key_id equals 0
+    See: https://core.telegram.org/mtproto/description#unencrypted-message
+    """
+    loaded_auth_key_id = load_long(buffer)
+
+    return loaded_auth_key_id.value == 0
+
+
+def is_encrypted_message(buffer: bytes) -> bool:
+    """
+    Encrypted message starts with auth_key_id that not equals to zero
+    See: https://core.telegram.org/mtproto/description#encrypted-message
+    """
+    return not is_unencrypted_message(buffer)
diff --git a/mtpylon/serialization/__init__.py b/mtpylon/serialization/__init__.py
index 6badf03..9916a9b 100644
--- a/mtpylon/serialization/__init__.py
+++ b/mtpylon/serialization/__init__.py
@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 from .schema import dump, load
 from .loaded import LoadedValue
+from .schema import CallableFunc
 
-__all__ = ['dump', 'load', 'LoadedValue']
+__all__ = ['dump', 'load', 'LoadedValue', 'CallableFunc']


Test patch:
diff --git a/tests/test_messages/__init__.py b/tests/test_messages/__init__.py
new file mode 100644
index 0000000..40a96af
--- /dev/null
+++ b/tests/test_messages/__init__.py
@@ -0,0 +1,1 @@
+# -*- coding: utf-8 -*-
diff --git a/tests/test_messages/test_unencrypted_message.py b/tests/test_messages/test_unencrypted_message.py
new file mode 100644
index 0000000..29c7e5f
--- /dev/null
+++ b/tests/test_messages/test_unencrypted_message.py
@@ -0,0 +1,92 @@
+# -*- coding: utf-8 -*-
+import pytest
+
+from mtpylon.messages.unencrypted_message import (
+    UnencryptedMessage,
+    unpack_message,
+    pack_message,
+)
+from mtpylon import long, int128
+from mtpylon.service_schema.functions import req_pq
+from mtpylon.service_schema.constructors import ResPQ
+from mtpylon.serialization.schema import CallableFunc
+
+
+encoded_req_pq_message = (
+    b'\x00\x00\x00\x00\x00\x00\x00\x00' +
+    b'\x4a\x96\x70\x27\xc4\x7a\xe5\x51' +
+    b'\x14\x00\x00\x00' +
+    b'\x78\x97\x46\x60' +
+    b'\x3e\x05\x49\x82\x8c\xca\x27\xe9\x66\xb3\x01\xa4\x8f\xec\xe2\xfc'
+)
+
+req_pq_message = UnencryptedMessage(
+    msg_id=long(0x51e57ac42770964a),
+    value=CallableFunc(
+        func=req_pq,
+        params={
+            'nonce': int128(0xfce2ec8fa401b366e927ca8c8249053e),
+        },
+    ),
+)
+
+encoded_res_pq = (
+    b'\x00\x00\x00\x00\x00\x00\x00\x00' +
+    b'\x01\x84\x67\xb9\xf4\x1d\x42\x60' +
+    b'\x40\x00\x00\x00' +
+    b'\x63\x24\x16\x05' +
+    b'\xb4\xbd\x9d\x5f\xd2\x98\x4f\x16\x04\x14\x4e\x16\x77\xd7\xce\x88' +
+    b'\x35\x5b\xea\x59\x40\xcc\x59\x3e\x7e\xc4\xc0\xbf\xc1\x23\x3d\x86' +
+    b'\x08\x13\x54\x65\x62\x4e\x61\x25\xfd\x00\x00\x00' +
+    b'\x15\xc4\xb5\x1c' +
+    b'\x01\x00\x00\x00' +
+    b'\x21\x6b\xe8\x6c\x02\x2b\xb4\xc3'
+)
+
+res_pq_message = UnencryptedMessage(
+    msg_id=long(0x60421df4b9678401),
+    value=ResPQ(
+        nonce=int128(0x88ced777164e1404164f98d25f9dbdb4),
+        server_nonce=int128(0x863d23c1bfc0c47e3e59cc4059ea5b35),
+        pq=b'\x13\x54\x65\x62\x4e\x61\x25\xfd',
+        server_public_key_fingerprints=[
+            long(0xc3b42b026ce86b21),
+        ],
+    ),
+)
+
+
+message_tests = [
+    pytest.param(
+        encoded_req_pq_message,
+        req_pq_message,
+        id='req_pq'
+    ),
+    pytest.param(
+        encoded_res_pq,
+        res_pq_message,
+        id='ResPQ'
+    ),
+]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'encoded_message,expected_message',
+    message_tests
+)
+async def test_unpack(encoded_message, expected_message):
+    message = await unpack_message(encoded_message)
+
+    assert message == expected_message
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'encoded_message,message',
+    message_tests
+)
+async def test_pack(encoded_message, message):
+    encoded = await pack_message(message)
+
+    assert encoded == encoded_message
diff --git a/tests/test_messages/test_utils.py b/tests/test_messages/test_utils.py
new file mode 100644
index 0000000..83163f0
--- /dev/null
+++ b/tests/test_messages/test_utils.py
@@ -0,0 +1,59 @@
+# -*- coding: utf-8 -*-
+import pytest
+
+from mtpylon.messages import is_encrypted_message, is_unencrypted_message
+
+
+unencrypted_message = (
+        b'\x00\x00\x00\x00\x00\x00\x00\x00' +
+        b'\x4a\x96\x70\x27\xc4\x7a\xe5\x51' +
+        b'\x14\x00\x00\x00' +
+        b'\x78\x97\x46\x60' +
+        b'\x3e\x05\x49\x82\x8c\xca\x27\xe9\x66\xb3\x01\xa4\x8f\xec\xe2\xfc'
+)
+
+encrypted_message = (
+        b'\x28\x43\x00\x00\x00\x00\x00\x00' +
+        b'\x4a\x96\x70\x27\xc4\x7a\xe5\x51' +
+        b'\x14\x00\x00\x00' +
+        b'\x78\x97\x46\x60' +
+        b'\x3e\x05\x49\x82\x8c\xca\x27\xe9\x66\xb3\x01\xa4\x8f\xec\xe2\xfc'
+)
+
+
+@pytest.mark.parametrize(
+    'dumped_message,value',
+    [
+        pytest.param(
+            unencrypted_message,
+            False,
+            id='unencrypted_message'
+        ),
+        pytest.param(
+            encrypted_message,
+            True,
+            id='encrypted_message'
+        ),
+    ]
+)
+def test_is_encrypted_message(dumped_message, value):
+    assert is_encrypted_message(dumped_message) == value
+
+
+@pytest.mark.parametrize(
+    'dumped_message,value',
+    [
+        pytest.param(
+            unencrypted_message,
+            True,
+            id='unencrypted_message'
+        ),
+        pytest.param(
+            encrypted_message,
+            False,
+            id='encrypted_message'
+        ),
+    ]
+)
+def test_is_unencrypted_message(dumped_message, value):
+    assert is_unencrypted_message(dumped_message) == value