Repository: Zapix/mtpylon (language: python)

Issue:
Transport obfuscation



Discussion / hints:

Fix patch:
diff --git a/mtpylon/transports/obfuscation.py b/mtpylon/transports/obfuscation.py
new file mode 100644
index 0000000..8f037ac
--- /dev/null
+++ b/mtpylon/transports/obfuscation.py
@@ -0,0 +1,57 @@
+# -*- coding: utf-8 -*-
+from typing import Tuple
+from dataclasses import dataclass
+
+import pyaes  # type: ignore
+
+
+@dataclass
+class Obfuscator:
+    """
+    Obfuscator for mtproto transport protocol.
+    Stores aes key/iv for client and server.
+    Provides `decrypt` method to decrypt client messages
+    Provides `encrypt` method to encrypt server messages
+    """
+    client_key: bytes
+    client_iv: bytes
+
+    server_key: bytes
+    server_iv: bytes
+
+    def __post_init__(self):
+        self._decipher = pyaes.AESModeOfOperationCTR(
+            self.client_key,
+            counter=pyaes.Counter(int.from_bytes(self.client_iv, 'big'))
+        )
+
+        self._cipher = pyaes.AESModeOfOperationCTR(
+            self.server_key,
+            counter=pyaes.Counter(int.from_bytes(self.server_iv, 'big'))
+        )
+
+    def decrypt(self, encrypted_message: bytes) -> bytes:
+        return self._decipher.decrypt(encrypted_message)
+
+    def encrypt(self, original_message: bytes) -> bytes:
+        return self._cipher.encrypt(original_message)
+
+
+def parse_header(header: bytes) -> Tuple[int, Obfuscator]:
+    """
+    Parse incoming header to retrieve client/server aes key iv and
+    transport protocol version that should be used.
+    Returns transport protocol tag and obfuscator
+    """
+    client_key = header[8:40]
+    client_iv = header[40:56]
+
+    reversed_header = header[::-1]
+    server_key = reversed_header[8:40]
+    server_iv = reversed_header[40:56]
+
+    obfuscator = Obfuscator(client_key, client_iv, server_key, server_iv)
+
+    encrypted_key = obfuscator.decrypt(header)
+
+    return int.from_bytes(encrypted_key[56:60], 'big'), obfuscator
diff --git a/requirements.txt b/requirements.txt
index 7cccf98..cea985d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -19,3 +19,4 @@ typed-ast==1.4.1
 typing-extensions==3.7.4.3
 rsa==4.7.2
 tgcrypto==1.2.2
+pyaes==1.6.1


Test patch:
diff --git a/tests/helpers.py b/tests/helpers.py
new file mode 100644
index 0000000..ebf8d01
--- /dev/null
+++ b/tests/helpers.py
@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+def hexstr_to_bytes(value: str) -> bytes:
+    return int(value, 16).to_bytes(len(value) // 2, 'big')
+
+
+def bytes_to_hexstr(value: bytes) -> str:
+    return '{0:0x}'.format(int.from_bytes(value, 'big'))
diff --git a/tests/test_transports/test_obfuscation.py b/tests/test_transports/test_obfuscation.py
new file mode 100644
index 0000000..a335dd9
--- /dev/null
+++ b/tests/test_transports/test_obfuscation.py
@@ -0,0 +1,76 @@
+# -*- coding: utf-8 -*-
+from mtpylon.transports.obfuscation import parse_header
+
+from ..helpers import hexstr_to_bytes
+
+
+init_buffer = hexstr_to_bytes(
+    'ac12ac1df7f68e2eacd4085138d294d7f9d71000469b9fbdc6cf0cbb50382a5d34a38c29ae4e14945dded705a65625e02e206e7a12f9e13aeeeeeeeece9fb887'  # noqa
+)
+
+reversed_init_buffer = hexstr_to_bytes(
+    '87b89fceeeeeeeee3ae1f9127a6e202ee02556a605d7de5d94144eae298ca3345d2a3850bb0ccfc6bd9f9b460010d7f9d794d2385108d4ac2e8ef6f71dac12ac'  # noqa
+)
+
+encrypted_init_buffer = hexstr_to_bytes(
+    '40a75bc91ea1eb6f367241fbfe013c2e56acc740dde85e8d7796af83763b552e567443c3780a6e96a5bd457b711c9a52963c217d99f8b8dc6bdb572fe929ec85'  # noqa
+)
+
+client_key = hexstr_to_bytes(
+    'acd4085138d294d7f9d71000469b9fbdc6cf0cbb50382a5d34a38c29ae4e1494'
+)
+
+client_iv = hexstr_to_bytes(
+    '5dded705a65625e02e206e7a12f9e13a'
+)
+
+server_key = hexstr_to_bytes(
+    '3ae1f9127a6e202ee02556a605d7de5d94144eae298ca3345d2a3850bb0ccfc6'
+
+)
+
+server_iv = hexstr_to_bytes(
+    'bd9f9b460010d7f9d794d2385108d4ac'
+)
+
+header = hexstr_to_bytes(
+    'ac12ac1df7f68e2eacd4085138d294d7f9d71000469b9fbdc6cf0cbb50382a5d34a38c29ae4e14945dded705a65625e02e206e7a12f9e13a6bdb572fe929ec85'  # noqa
+)
+
+client_original_msg = hexstr_to_bytes(
+    '280000000000000000000000040000005e595260140000007897466054b3ae2f3172ad489eef5aa8271f730e'  # noqa
+)
+
+client_encrypted_msg = hexstr_to_bytes(
+    'e7dbc655fbf3234c5f265c50c06482bd55ac73a906fdfee24b946ea3d268e42e11a1fccac23884200ffb7e67'  # noqa
+)
+
+server_original_message = hexstr_to_bytes(
+    '54000000000000000000000001f442f55e595260400000006324160554b3ae2f3172ad489eef5aa8271f730e857a0d1dc11006c19103166c8b21d22b082690282a15c9613d00000015c4b51c01000000216be86c022bb4c3'  # noqa
+)
+
+server_encrypted_msg = hexstr_to_bytes(
+    'fb4209d8ae8eb555348b62e53c0df46e0bb3fc0dd8b8c676c8abd9200e5546d61997e1e8473bb39304afa06ab84a4cee580541bd99bdb602737e4d358ac5c609451b2234720d93120a37b9c0b460e626d27d9da196c4690c'  # noqa
+)
+
+
+def test_parsed_header():
+    transport_tag, obfuscator = parse_header(header)
+
+    assert transport_tag == 0xeeeeeeee
+    assert obfuscator.client_key == client_key
+    assert obfuscator.client_iv == client_iv
+    assert obfuscator.server_key == server_key
+    assert obfuscator.server_iv == server_iv
+
+
+def test_decrypt_client_message():
+    transport_tag, obfuscator = parse_header(header)
+
+    assert obfuscator.decrypt(client_encrypted_msg) == client_original_msg
+
+
+def test_encrypt_server_message():
+    transport_tag, obfuscator = parse_header(header)
+
+    assert obfuscator.encrypt(server_original_message) == server_encrypted_msg