From 260e006abe9f2bdad5907c4b3e6e801f0df7e870 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Fri, 13 Oct 2017 11:28:19 +0200
Subject: [PATCH] Initial attempt at Perfect Forward Secrecy

---
 telethon/crypto/auth_key.py       |  6 +++
 telethon/helpers.py               | 15 ++++++-
 telethon/network/authenticator.py | 72 +++++++++++++++++++++++++++++--
 telethon/telegram_bare_client.py  |  2 +-
 telethon/tl/session.py            |  2 +-
 telethon_generator/scheme.tl      |  2 +
 6 files changed, 91 insertions(+), 8 deletions(-)

diff --git a/telethon/crypto/auth_key.py b/telethon/crypto/auth_key.py
index 17a7f8ca..fc8534d0 100644
--- a/telethon/crypto/auth_key.py
+++ b/telethon/crypto/auth_key.py
@@ -21,3 +21,9 @@ class AuthKey:
         new_nonce = new_nonce.to_bytes(32, 'little', signed=True)
         data = new_nonce + struct.pack('<BQ', number, self.aux_hash)
         return utils.calc_msg_key(data)
+
+
+class TempAuthKey(AuthKey):
+    def __init__(self, data, expires_at):
+        super().__init__(data)
+        self.expires_at = expires_at
diff --git a/telethon/helpers.py b/telethon/helpers.py
index 3c9af2cb..a990148d 100644
--- a/telethon/helpers.py
+++ b/telethon/helpers.py
@@ -1,13 +1,14 @@
 """Various helpers not related to the Telegram API itself"""
+import struct
 from hashlib import sha1, sha256
 import os
 
 # region Multiple utilities
 
 
-def generate_random_long(signed=True):
+def random_long(signed=True):
     """Generates a random long integer (8 bytes), which is optionally signed"""
-    return int.from_bytes(os.urandom(8), signed=signed, byteorder='little')
+    return struct.unpack('<q' if signed else '<Q', os.urandom(8))[0]
 
 
 def ensure_parent_dir_exists(file_path):
@@ -72,4 +73,14 @@ def get_password_hash(pw, current_salt):
     pw_hash = current_salt + data + current_salt
     return sha256(pw_hash).digest()
 
+
+def reinterpret(value, fmt, endian='<'):
+    """Reinterprets the given value from its original format to its new
+       format, for instance, reinterpret(-1, 'qQ') -> -1 as unsigned.
+    """
+    return struct.unpack(
+        endian + fmt[1], struct.pack(endian + fmt[0], value)
+    )[0]
+
+
 # endregion
diff --git a/telethon/network/authenticator.py b/telethon/network/authenticator.py
index 78df5d87..9bd6be1f 100644
--- a/telethon/network/authenticator.py
+++ b/telethon/network/authenticator.py
@@ -1,20 +1,25 @@
 import os
+import struct
 import time
+from datetime import datetime, timedelta
 from hashlib import sha1
 
-from ..tl.types import (
-    ResPQ, PQInnerData, ServerDHParamsFail, ServerDHParamsOk,
-    ServerDHInnerData, ClientDHInnerData, DhGenOk, DhGenRetry, DhGenFail
-)
 from .. import helpers as utils
 from ..crypto import AES, AuthKey, Factorization
 from ..crypto import rsa
 from ..errors import SecurityError
 from ..extensions import BinaryReader
 from ..network import MtProtoPlainSender
+from ..tl import TLMessage, TLObject
 from ..tl.functions import (
     ReqPqRequest, ReqDHParamsRequest, SetClientDHParamsRequest
 )
+from ..tl.functions.auth import BindTempAuthKeyRequest
+from ..tl.types import (
+    ResPQ, PQInnerData, ServerDHParamsFail, ServerDHParamsOk,
+    ServerDHInnerData, ClientDHInnerData, DhGenOk, DhGenRetry, DhGenFail,
+    BindAuthKeyInner
+)
 
 
 def do_authentication(connection, retries=5):
@@ -190,6 +195,65 @@ def _do_authentication(connection):
         raise NotImplementedError('DH Gen unknown: {}'.format(dh_gen))
 
 
+def bind_temp_auth_key(temp_auth_key, auth_key, sender):
+    """
+    :param TempAuthKey temp_auth_key:
+    :param AuthKey auth_key:
+    :param MtProtoSender sender:
+    """
+    assert sender.session.auth_key.key_id == temp_auth_key.key_id
+    bind_inner = BindAuthKeyInner(
+        nonce=utils.random_long(),
+        temp_auth_key_id=utils.reinterpret(temp_auth_key.key_id, 'Qq'),
+        perm_auth_key_id=utils.reinterpret(sender.session.auth_key.key_id, 'Qq'),
+        temp_session_id=utils.reinterpret(sender.session.id, 'Qq'),
+        expires_at=datetime.now() + timedelta(days=1)
+    )
+    inner_message = TLMessage(sender.session, bind_inner)
+    inner_message.seq_no = 0
+
+    # TODO Copy paste from MtProtoSender._send_message
+    plain_text = (
+        os.urandom(16)  # random unsigned salt + random unsigned id
+        + inner_message.to_bytes()
+    )
+    msg_key = utils.calc_msg_key(plain_text)
+    key_id = struct.pack('<Q', auth_key.key_id)
+    key, iv = utils.calc_key(auth_key.key, msg_key, client=True)
+    cipher_text = AES.encrypt_ige(plain_text, key, iv)
+    encrypted_inner_message = key_id + msg_key + cipher_text
+
+    # Now perform the actual request
+    request = BindTempAuthKeyRequest(
+        perm_auth_key_id=utils.reinterpret(auth_key.key_id, 'Qq'),
+        nonce=bind_inner.nonce,
+        expires_at=bind_inner.expires_at,
+        encrypted_message=TLObject.serialize_bytes(encrypted_inner_message)
+    )
+
+    message = TLMessage(sender.session, request)
+    message.msg_id = inner_message.msg_id
+
+    # TODO Copy paste from MtProtoSender._send_message
+    plain_text = \
+        struct.pack('<QQ', sender.session.salt, sender.session.id) \
+        + message.to_bytes()
+
+    msg_key = utils.calc_msg_key(plain_text)
+    key_id = struct.pack('<Q', temp_auth_key.key_id)
+    key, iv = utils.calc_key(temp_auth_key.key, msg_key, client=True)
+    cipher_text = AES.encrypt_ige(plain_text, key, iv)
+    encrypted_message = key_id + msg_key + cipher_text
+
+    sender._pending_receive[message.msg_id] = message
+    sender.connection.send(encrypted_message)
+    while not request.confirm_received.is_set():
+        sender.receive(update_state=None)
+
+    print('Success?')
+    print(request.result.stringify())
+
+
 def get_int(byte_array, signed=True):
     """Gets the specified integer from its byte array.
        This should be used by the authenticator,
diff --git a/telethon/telegram_bare_client.py b/telethon/telegram_bare_client.py
index e8e8b7fb..bfe22d25 100644
--- a/telethon/telegram_bare_client.py
+++ b/telethon/telegram_bare_client.py
@@ -601,7 +601,7 @@ class TelegramBareClient:
         is_large = file_size > 10 * 1024 * 1024
         part_count = (file_size + part_size - 1) // part_size
 
-        file_id = utils.generate_random_long()
+        file_id = utils.random_long()
         hash_md5 = md5()
 
         stream = open(file, 'rb') if isinstance(file, str) else BytesIO(file)
diff --git a/telethon/tl/session.py b/telethon/tl/session.py
index f597048f..9f0f7e57 100644
--- a/telethon/tl/session.py
+++ b/telethon/tl/session.py
@@ -58,7 +58,7 @@ class Session:
         self._msg_id_lock = Lock()
         self._save_lock = Lock()
 
-        self.id = helpers.generate_random_long(signed=False)
+        self.id = helpers.random_long(signed=False)
         self._sequence = 0
         self.time_offset = 0
         self._last_msg_id = 0  # Long
diff --git a/telethon_generator/scheme.tl b/telethon_generator/scheme.tl
index 5e949239..a87cd38d 100644
--- a/telethon_generator/scheme.tl
+++ b/telethon_generator/scheme.tl
@@ -51,6 +51,8 @@ destroy_auth_key_ok#f660e1d4 = DestroyAuthKeyRes;
 destroy_auth_key_none#0a9f2259 = DestroyAuthKeyRes;
 destroy_auth_key_fail#ea109b13 = DestroyAuthKeyRes;
 
+bind_auth_key_inner#75a3f765 nonce:long temp_auth_key_id:long perm_auth_key_id:long temp_session_id:long expires_at:int = BindAuthKeyInner;
+
 ---functions---
 
 req_pq#60469778 nonce:int128 = ResPQ;