From 49ca5c00c76f435e03a48509a0d930e16333d45c Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Wed, 7 Jun 2017 11:48:54 +0100 Subject: [PATCH] Create a more upgrade-resistant JSON-based Session class --- telethon/telegram_client.py | 5 +- telethon/tl/__init__.py | 2 +- telethon/tl/session.py | 114 ++++++++++++++++++++++++++++++++++++ 3 files changed, 118 insertions(+), 3 deletions(-) diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py index d1dab1d1..01f6f44e 100644 --- a/telethon/telegram_client.py +++ b/telethon/telegram_client.py @@ -16,7 +16,7 @@ from .network import authenticator, MtProtoSender, TcpTransport from .parser.markdown_parser import parse_message_entities # For sending and receiving requests -from .tl import MTProtoRequest, Session +from .tl import MTProtoRequest, Session, JsonSession from .tl.all_tlobjects import layer from .tl.functions import (InitConnectionRequest, InvokeWithLayerRequest, PingRequest) @@ -85,8 +85,9 @@ class TelegramClient: self.api_hash = api_hash # Determine what session object we have + # TODO JsonSession until migration is complete (by v1.0) if isinstance(session, str) or session is None: - self.session = Session.try_load_or_create_new(session) + self.session = JsonSession.try_load_or_create_new(session) elif isinstance(session, Session): self.session = session else: diff --git a/telethon/tl/__init__.py b/telethon/tl/__init__.py index 0fa833da..8f58f5c2 100644 --- a/telethon/tl/__init__.py +++ b/telethon/tl/__init__.py @@ -1,2 +1,2 @@ from .mtproto_request import MTProtoRequest -from .session import Session +from .session import Session, JsonSession diff --git a/telethon/tl/session.py b/telethon/tl/session.py index d72b326f..14140648 100644 --- a/telethon/tl/session.py +++ b/telethon/tl/session.py @@ -1,7 +1,9 @@ +import json import os import pickle import random import time +from base64 import b64encode, b64decode from os.path import isfile as file_exists from .. import helpers as utils @@ -72,3 +74,115 @@ class Session: now = int(time.time()) correct = correct_msg_id >> 32 self.time_offset = correct - now + + +# Until migration is complete, we need the original 'Session' class +# for Pickle to keep working. TODO Replace 'Session' by 'JsonSession' by v1.0 +class JsonSession: + """This session contains the required information to login into your + Telegram account. NEVER give the saved JSON file to anyone, since + they would gain instant access to all your messages and contacts. + + If you think the session has been compromised, close all the sessions + through an official Telegram client to revoke the authorization. + """ + def __init__(self, session_user_id): + self.session_user_id = session_user_id + self.server_address = '91.108.56.165' + self.port = 443 + self.auth_key = None + self.id = utils.generate_random_long(signed=False) + self.sequence = 0 + self.salt = 0 # Unsigned long + self.time_offset = 0 + self.last_message_id = 0 # Long + + def save(self): + """Saves the current session object as session_user_id.session""" + if self.session_user_id: + with open('{}.session'.format(self.session_user_id), 'w') as file: + json.dump({ + 'id': self.id, + 'port': self.port, + 'salt': self.salt, + 'sequence': self.sequence, + 'time_offset': self.time_offset, + 'server_address': self.server_address, + 'auth_key_data': + b64encode(self.auth_key.key).decode('ascii') + }, file) + + def delete(self): + """Deletes the current session file""" + try: + os.remove('{}.session'.format(self.session_user_id)) + return True + except OSError: + return False + + @staticmethod + def try_load_or_create_new(session_user_id): + """Loads a saved session_user_id.session or creates a new one. + If session_user_id=None, later .save()'s will have no effect. + """ + if session_user_id is None: + return JsonSession(None) + else: + path = '{}.session'.format(session_user_id) + result = JsonSession(session_user_id) + if not file_exists(path): + return result + + try: + with open(path, 'r') as file: + data = json.load(file) + result.id = data['id'] + result.port = data['port'] + result.salt = data['salt'] + result.sequence = data['sequence'] + result.time_offset = data['time_offset'] + result.server_address = data['server_address'] + + # FIXME We need to import the AuthKey here or otherwise + # we get cyclic dependencies. + from ..crypto import AuthKey + key = b64decode(data['auth_key_data']) + result.auth_key = AuthKey(data=key) + + except (json.decoder.JSONDecodeError, UnicodeDecodeError): + # TODO Backwards-compatibility code + old = Session.try_load_or_create_new(session_user_id) + result.id = old.id + result.port = old.port + result.salt = old.salt + result.sequence = old.sequence + result.time_offset = old.time_offset + result.server_address = old.server_address + result.auth_key = old.auth_key + result.save() + + return result + + def get_new_msg_id(self): + """Generates a new unique message ID based on the current + time (in ms) since epoch""" + # Refer to mtproto_plain_sender.py for the original method, + ms_time = int(time.time() * 1000) + new_msg_id = (((ms_time // 1000 + self.time_offset) << 32) + | # "must approximately equal unix time*2^32" + ((ms_time % 1000) << 22) + | # "approximate moment in time the message was created" + random.randint(0, 524288) + << 2) # "message identifiers are divisible by 4" + + if self.last_message_id >= new_msg_id: + new_msg_id = self.last_message_id + 4 + + self.last_message_id = new_msg_id + return new_msg_id + + def update_time_offset(self, correct_msg_id): + """Updates the time offset based on a known correct message ID""" + now = int(time.time()) + correct = correct_msg_id >> 32 + self.time_offset = correct - now