Xtao-Labs/MTPyroger
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Reorganize Session to make use of the MTProto module

Browse Source
This commit is contained in:
Dan 2018-06-14 03:22:52 +02:00
parent 75121c9c57
commit 11ddf5f99d
1 changed files with 52 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

91
pyrogram/session/session.py
View File

@ -19,20 +19,18 @@
import asyncio import asyncio
import logging import logging
import platform import platform
import threading
from datetime import datetime, timedelta from datetime import datetime, timedelta
from hashlib import sha1, sha256 from hashlib import sha1
from io import BytesIO from io import BytesIO
from os import urandom
import pyrogram import pyrogram
from pyrogram import __copyright__, __license__, __version__ from pyrogram import __copyright__, __license__, __version__
from pyrogram.api import functions, types, core from pyrogram.api import functions, types, core
from pyrogram.api.all import layer from pyrogram.api.all import layer
from pyrogram.api.core import Message, Object, MsgContainer, Long, FutureSalt, Int from pyrogram.api.core import Object, MsgContainer, Long, FutureSalt, Int
from pyrogram.api.errors import Error, InternalServerError from pyrogram.api.errors import Error, InternalServerError
from pyrogram.connection import Connection from pyrogram.connection import Connection
from pyrogram.crypto import AES, KDF from pyrogram.crypto import MTProto
from .internals import MsgId, MsgFactory, DataCenter from .internals import MsgId, MsgFactory, DataCenter
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
@ -58,7 +56,6 @@ class Session:
platform.release() platform.release()
) )
INITIAL_SALT = 0x616e67656c696361
NET_WORKERS = 1 NET_WORKERS = 1
WAIT_TIMEOUT = 15 WAIT_TIMEOUT = 15
MAX_RETRIES = 5 MAX_RETRIES = 5
@ -137,7 +134,7 @@ class Session:
self.net_worker_task = asyncio.ensure_future(self.net_worker()) self.net_worker_task = asyncio.ensure_future(self.net_worker())
self.recv_task = asyncio.ensure_future(self.recv()) self.recv_task = asyncio.ensure_future(self.recv())
self.current_salt = FutureSalt(0, 0, self.INITIAL_SALT) self.current_salt = FutureSalt(0, 0, MTProto.INITIAL_SALT)
self.current_salt = FutureSalt(0, 0, (await self._send(functions.Ping(0))).new_server_salt) self.current_salt = FutureSalt(0, 0, (await self._send(functions.Ping(0))).new_server_salt)
self.current_salt = (await self._send(functions.GetFutureSalts(1))).salts[0] self.current_salt = (await self._send(functions.GetFutureSalts(1))).salts[0]
@ -215,36 +212,40 @@ class Session:
data = Long(self.current_salt.salt) + self.session_id + message.write() data = Long(self.current_salt.salt) + self.session_id + message.write()
padding = urandom(-(len(data) + 12) % 16 + 12) padding = urandom(-(len(data) + 12) % 16 + 12)
# 88 = 88 + 0 (outgoing message) # def pack(self, message: Message):
msg_key_large = sha256(self.auth_key[88: 88 + 32] + data + padding).digest() # data = Long(self.current_salt.salt) + self.session_id + message.write()
msg_key = msg_key_large[8:24] # padding = urandom(-(len(data) + 12) % 16 + 12)
aes_key, aes_iv = KDF(self.auth_key, msg_key, True) #
# # 88 = 88 + 0 (outgoing message)
return self.auth_key_id + msg_key + AES.ige256_encrypt(data + padding, aes_key, aes_iv) # msg_key_large = sha256(self.auth_key[88: 88 + 32] + data + padding).digest()
# msg_key = msg_key_large[8:24]
def unpack(self, b: BytesIO) -> Message: # aes_key, aes_iv = KDF(self.auth_key, msg_key, True)
assert b.read(8) == self.auth_key_id, b.getvalue() #
# return self.auth_key_id + msg_key + AES.ige256_encrypt(data + padding, aes_key, aes_iv)
msg_key = b.read(16) #
aes_key, aes_iv = KDF(self.auth_key, msg_key, False) # def unpack(self, b: BytesIO) -> Message:
data = BytesIO(AES.ige256_decrypt(b.read(), aes_key, aes_iv)) # assert b.read(8) == self.auth_key_id, b.getvalue()
data.read(8) #
# msg_key = b.read(16)
# https://core.telegram.org/mtproto/security_guidelines#checking-session-id # aes_key, aes_iv = KDF(self.auth_key, msg_key, False)
assert data.read(8) == self.session_id # data = BytesIO(AES.ige256_decrypt(b.read(), aes_key, aes_iv))
# data.read(8)
message = Message.read(data) #
# # https://core.telegram.org/mtproto/security_guidelines#checking-session-id
# https://core.telegram.org/mtproto/security_guidelines#checking-sha256-hash-value-of-msg-key # assert data.read(8) == self.session_id
# https://core.telegram.org/mtproto/security_guidelines#checking-message-length #
# 96 = 88 + 8 (incoming message) # message = Message.read(data)
assert msg_key == sha256(self.auth_key[96:96 + 32] + data.getvalue()).digest()[8:24] #
# # https://core.telegram.org/mtproto/security_guidelines#checking-sha256-hash-value-of-msg-key
# https://core.telegram.org/mtproto/security_guidelines#checking-msg-id # # https://core.telegram.org/mtproto/security_guidelines#checking-message-length
# TODO: check for lower msg_ids # # 96 = 88 + 8 (incoming message)
assert message.msg_id % 2 != 0 # assert msg_key == sha256(self.auth_key[96:96 + 32] + data.getvalue()).digest()[8:24]
#
return message # # https://core.telegram.org/mtproto/security_guidelines#checking-msg-id
# # TODO: check for lower msg_ids
# assert message.msg_id % 2 != 0
#
# return message
async def net_worker(self): async def net_worker(self):
name = threading.current_thread().name name = threading.current_thread().name
@ -257,7 +258,13 @@ class Session:
break break
try: try:
data = self.unpack(BytesIO(packet)) data = MTProto.unpack(
BytesIO(packet),
self.current_salt.salt,
self.session_id,
self.auth_key,
self.auth_key_id
)
messages = ( messages = (
data.body.messages data.body.messages
@ -391,7 +398,13 @@ class Session:
if wait_response: if wait_response:
self.results[msg_id] = Result() self.results[msg_id] = Result()
payload = self.pack(message) payload = MTProto.pack(
message,
self.current_salt.salt,
self.session_id,
self.auth_key,
self.auth_key_id
)
try: try:
await self.connection.send(payload) await self.connection.send(payload)