From 9420e152830dd32ecef3a2fefdb25d9c682ca435 Mon Sep 17 00:00:00 2001 From: Lonami Date: Mon, 12 Sep 2016 19:32:16 +0200 Subject: [PATCH] Gave more power to the TelegramClients and bug fixes Fixed uploads for large files on TcpClient Added more RPCError's for handling invalid phone code Added more media handlers: now you're also able to both send and download documents! The InteractiveTelegramClient now supports working with media aswell --- .gitignore | 2 +- README.md | 3 +- errors.py | 6 +- interactive_telegram_client.py | 89 +++++++++++--------- network/tcp_client.py | 14 +++- telegram_client.py | 143 +++++++++++++++++++++++++++++---- utils/helpers.py | 9 +++ 7 files changed, 208 insertions(+), 58 deletions(-) diff --git a/.gitignore b/.gitignore index b08a8a02..2a8a7662 100755 --- a/.gitignore +++ b/.gitignore @@ -5,7 +5,7 @@ tl/all_tlobjects.py # User session *.session -*.jpg +usermedia/ api/settings # Byte-compiled / optimized / DLL files diff --git a/README.md b/README.md index 3daa78a2..0349c172 100755 --- a/README.md +++ b/README.md @@ -49,7 +49,8 @@ Then fill the file with the corresponding values (your `api_id`, `api_hash` and ## Running Telethon First of all, you need to run the `tl_generator.py` by issuing `python3 tl_generator.py`. This will generate all the -TLObjects from the given `scheme.tl` file. When it's done, you can run `python3 main.py` to start the interactive example. +TLObjects from the given `scheme.tl` file. When it's done, you can run `python3 interactive_telegram_client.py` to +start the interactive example. ## Advanced uses ### Using more than just `TelegramClient` diff --git a/errors.py b/errors.py index 56e90075..e2197542 100644 --- a/errors.py +++ b/errors.py @@ -93,9 +93,11 @@ class RPCError(Exception): 'PHONE_NUMBER_INVALID': 'The phone number is invalid.', - 'PHONE_CODE_HASH_EMPTY': 'phone_code_hash is missing.', + 'PHONE_CODE_HASH_EMPTY': 'The phone code hash is missing.', - 'PHONE_CODE_EMPTY': 'phone_code is missing.', + 'PHONE_CODE_EMPTY': 'The phone code is missing.', + + 'PHONE_CODE_INVALID': 'The phone code entered was invalid.', 'PHONE_CODE_EXPIRED': 'The confirmation code has expired.', diff --git a/interactive_telegram_client.py b/interactive_telegram_client.py index 3d4fbe6a..30ece0bd 100644 --- a/interactive_telegram_client.py +++ b/interactive_telegram_client.py @@ -1,5 +1,4 @@ import tl_generator -from tl.types import MessageMediaPhoto from tl.types import UpdateShortChatMessage from tl.types import UpdateShortMessage @@ -12,12 +11,8 @@ else: from telegram_client import TelegramClient from utils.helpers import load_settings - -# For pretty printing, thanks to http://stackoverflow.com/a/37501797/4759433 -import sys -import readline -from time import sleep import shutil +import traceback # Get the (current) number of lines in the terminal cols, rows = shutil.get_terminal_size() @@ -51,8 +46,10 @@ class InteractiveTelegramClient(TelegramClient): print('First run. Sending code request...') self.send_code_request(user_phone) - code = input('Enter the code you just received: ') - self.make_auth(user_phone, code) + code_ok = False + while not code_ok: + code = input('Enter the code you just received: ') + code_ok = self.make_auth(user_phone, code) def run(self): # Listen for updates @@ -95,9 +92,11 @@ class InteractiveTelegramClient(TelegramClient): print_title('Chat with "{}"'.format(display)) print('Available commands:'.format(display)) print(' !q: Quits the current chat.') + print(' !Q: Quits the current chat and exits.') print(' !h: prints the latest messages (message History) of the chat.') - print(' !p : sends a Photo located at the given path') - print(' !d : Downloads the given message media (if any)') + print(' !p : sends a Photo located at the given path.') + print(' !f : sends a File document located at the given path.') + print(' !d : Downloads the given message media (if any).') # And start a while loop to chat while True: @@ -105,6 +104,8 @@ class InteractiveTelegramClient(TelegramClient): # Quit if msg == '!q': break + elif msg == '!Q': + return # History elif msg == '!h': @@ -130,40 +131,56 @@ class InteractiveTelegramClient(TelegramClient): # Send photo elif msg.startswith('!p '): - file_path = msg[len('!p '):] # Slice the message to get the path + # Slice the message to get the path + self.send_photo(path=msg[len('!p '):], peer=input_peer) - print('Uploading {}...'.format(file_path)) - input_file = self.upload_file(file_path) - - # After we have the handle to the uploaded file, send it to our peer - self.send_photo_file(input_file, input_peer) - print('Media sent!') + # Send file (document) + elif msg.startswith('!f '): + # Slice the message to get the path + self.send_document(path=msg[len('!f '):], peer=input_peer) # Download media elif msg.startswith('!d '): - msg_media_id = msg[len('!d '):] # Slice the message to get message ID - try: - # The user may have entered a non-integer string! - msg_media_id = int(msg_media_id) - - # Search the message ID and ensure the media is a Photo - for msg in self.found_media: - if (msg.id == msg_media_id and - type(msg.media) == MessageMediaPhoto): - - # Retrieve the output and download the photo - output = '{}.jpg'.format(str(msg_media_id)) - print('Downloading to {}...'.format(output)) - self.download_photo(msg.media, file_path=output) - print('Photo downloaded to {}!'.format(output)) - - except ValueError: - print('Invalid media ID given!') + # Slice the message to get message ID + self.download_media(msg[len('!d '):]) # Send chat message (if any) elif msg: self.send_message(input_peer, msg, markdown=True, no_web_page=True) + def send_photo(self, path, peer): + print('Uploading {}...'.format(path)) + input_file = self.upload_file(path) + + # After we have the handle to the uploaded file, send it to our peer + self.send_photo_file(input_file, peer) + print('Photo sent!') + + def send_document(self, path, peer): + print('Uploading {}...'.format(path)) + input_file = self.upload_file(path) + + # After we have the handle to the uploaded file, send it to our peer + self.send_document_file(input_file, peer) + print('Document sent!') + + def download_media(self, media_id): + try: + # The user may have entered a non-integer string! + msg_media_id = int(media_id) + + # Search the message ID + for msg in self.found_media: + if msg.id == msg_media_id: + # Let the output be the message ID + output = str('usermedia/{}'.format(msg_media_id)) + print('Downloading media with name {}...'.format(output)) + output = self.download_msg_media(msg.media, file_path=output) + print('Media downloaded to {}!'.format(output)) + + except ValueError: + print('Invalid media ID given!') + @staticmethod def update_handler(update_object): if type(update_object) is UpdateShortMessage: @@ -189,7 +206,7 @@ if __name__ == '__main__': client.run() except Exception as e: - print('Unexpected error ({}), will not continue: {}'.format(type(e), e)) + print('Unexpected error ({}): {} at\n{}', type(e), e, traceback.format_exc()) finally: print_title('Exit') diff --git a/network/tcp_client.py b/network/tcp_client.py index 571163ab..14b3043e 100755 --- a/network/tcp_client.py +++ b/network/tcp_client.py @@ -34,7 +34,18 @@ class TcpClient: # Ensure that only one thread can send data at once with self.lock: - self.socket.sendall(data) + # Do not use .sendall: + # "on error, an exception is raised, and there is no way to + # determine how much data, if any, was successfully sent." + while data: + try: + sent = self.socket.send(data) + data = data[sent:] + except BlockingIOError as e: + if 'Errno 11' in str(e): # Error #11: Resource temporary unavailable + time.sleep(0.1) # Sleep a bit waiting for the resource to be available + else: + raise e def read(self, buffer_size): """Reads (receives) the specified bytes from the connected peer""" @@ -65,6 +76,7 @@ class TcpClient: # If everything went fine, return the read bytes return writer.get_bytes() + def cancel_read(self): """Cancels the read operation IF it hasn't yet started, raising a ReadCancelledError""" diff --git a/telegram_client.py b/telegram_client.py index a0cd9745..32d5cb95 100644 --- a/telegram_client.py +++ b/telegram_client.py @@ -2,6 +2,7 @@ import platform from datetime import datetime from hashlib import md5 from os import path +from mimetypes import guess_extension, guess_type import utils import network.authenticator @@ -18,7 +19,9 @@ from tl import Session from tl.types import \ PeerUser, PeerChat, PeerChannel, \ InputPeerUser, InputPeerChat, InputPeerChannel, InputPeerEmpty, \ - InputFile, InputFileLocation, InputMediaUploadedPhoto + InputFile, InputFileLocation, InputMediaUploadedPhoto, InputMediaUploadedDocument, \ + MessageMediaContact, MessageMediaDocument, MessageMediaPhoto, \ + DocumentAttributeAudio, DocumentAttributeFilename, InputDocumentFileLocation from tl.functions import InvokeWithLayerRequest, InitConnectionRequest from tl.functions.help import GetConfigRequest @@ -143,10 +146,16 @@ class TelegramClient: if phone_number not in self.phone_code_hashes: raise ValueError('Please make sure you have called send_code_request first.') - # TODO Handle invalid code - request = SignInRequest(phone_number, self.phone_code_hashes[phone_number], code) - self.sender.send(request) - self.sender.receive(request) + try: + request = SignInRequest(phone_number, self.phone_code_hashes[phone_number], code) + self.sender.send(request) + self.sender.receive(request) + except RPCError as error: + if error.message.startswith('PHONE_CODE_'): + print(error) + return False + else: + raise error # Result is an Auth.Authorization TLObject self.session.user = request.result.user @@ -155,7 +164,7 @@ class TelegramClient: # Now that we're authorized, we can listen for incoming updates self.sender.set_listen_for_updates(True) - return self.session.user + return True # endregion @@ -229,14 +238,13 @@ class TelegramClient: return total_messages, result.messages, users - # endregion - # region Uploading/downloading media requests - # TODO Handle media downloading/uploading in a different session? # "It is recommended that large queries (upload.getFile, upload.saveFilePart) # be handled through a separate session and a separate connection" + # region Uploading media requests + def upload_file(self, file_path, part_size_kb=64, file_name=None): """Uploads the specified media with the given chunk (part) size, in KB. If no custom file name is specified, the real file name will be used. @@ -262,6 +270,8 @@ class TelegramClient: if not part: break + print('I read {} out of {}'.format(len(part), part_size)) + # Invoke the file upload and increment both the part index and MD5 checksum result = self.invoke(SaveFilePartRequest(file_id, part_index, part)) if result: @@ -286,37 +296,136 @@ class TelegramClient: self.send_media_file( InputMediaUploadedPhoto(input_file, caption), input_peer) + def send_document_file(self, input_file, input_peer, caption=''): + """Sends a previously uploaded input_file + (which should be a document) to an input_peer""" + + # Determine mime-type and attributes + # Take the first element by using [0] since it returns a tuple + mime_type = guess_type(input_file.name)[0] + attributes = [ + DocumentAttributeFilename(input_file.name) + # TODO If the input file is an audio, find out: + # Performer and song title and add DocumentAttributeAudio + ] + self.send_media_file(InputMediaUploadedDocument(file=input_file, + mime_type=mime_type, + attributes=attributes, + caption=caption), input_peer) + def send_media_file(self, input_media, input_peer): """Sends any input_media (contact, document, photo...) to an input_peer""" self.invoke(SendMediaRequest(peer=input_peer, media=input_media, random_id=utils.generate_random_long())) - def download_photo(self, message_media_photo, file_path): - """Downloads a message_media_photo largest size into the desired file_path""" + # endregion + + # region Downloading media requests + + def download_msg_media(self, message_media, file_path, add_extension=True): + """Downloads the given MessageMedia (Photo, Document or Contact) + into the desired file_path, optionally finding its extension automatically""" + if type(message_media) == MessageMediaPhoto: + return self.download_photo(message_media, file_path, add_extension) + + elif type(message_media) == MessageMediaDocument: + return self.download_document(message_media, file_path, add_extension) + + elif type(message_media) == MessageMediaContact: + return self.download_contact(message_media, file_path, add_extension) + + def download_photo(self, message_media_photo, file_path, add_extension=False): + """Downloads MessageMediaPhoto's largest size into the desired + file_path, optionally finding its extension automatically""" # Determine the photo and its largest size photo = message_media_photo.photo largest_size = photo.sizes[-1].location - # Download the media with the largest size input file location - self.download_media(InputFileLocation(volume_id=largest_size.volume_id, - local_id=largest_size.local_id, - secret=largest_size.secret), file_path) + # Photos are always compressed into a .jpg by Telegram + if add_extension: + file_path += '.jpg' - def download_media(self, input_file_location, file_path, part_size_kb=64): + # Download the media with the largest size input file location + self.download_file_loc(InputFileLocation(volume_id=largest_size.volume_id, + local_id=largest_size.local_id, + secret=largest_size.secret), file_path) + return file_path + + def download_document(self, message_media_document, file_path=None, add_extension=True): + """Downloads the given MessageMediaDocument into the desired + file_path, optionally finding its extension automatically. + If no file_path is given, it will _try_ to be guessed from the document""" + document = message_media_document.document + + # If no file path was given, try to guess it from the attributes + if file_path is None: + for attr in document.attributes: + if type(attr) == DocumentAttributeFilename: + file_path = attr.file_name + break # This attribute has higher preference + + elif type(attr) == DocumentAttributeAudio: + file_path = '{} - {}'.format(attr.performer, attr.title) + + if file_path is None: + print('Could not determine a filename for the document') + + # Guess the extension based on the mime_type + if add_extension: + ext = guess_extension(document.mime_type) + if ext is not None: + file_path += ext + + self.download_file_loc(InputDocumentFileLocation(id=document.id, + access_hash=document.access_hash, + version=document.version), file_path) + + return file_path + + @staticmethod + def download_contact(message_media_contact, file_path, add_extension=True): + """Downloads a media contact using the vCard 4.0 format""" + + first_name = message_media_contact.first_name + last_name = message_media_contact.last_name + phone_number = message_media_contact.phone_number + + # The only way we can save a contact in an understandable + # way by phones is by using the .vCard format + if add_extension: + file_path += '.vcard' + + # Ensure that we'll be able to download the contact + utils.ensure_parent_dir_exists(file_path) + + with open(file_path, 'w', encoding='utf-8') as file: + file.write('BEGIN:VCARD\n') + file.write('VERSION:4.0\n') + file.write('N:{};{};;;\n'.format(first_name, last_name if last_name else '')) + file.write('FN:{}\n'.format(' '.join((first_name, last_name)))) + file.write('TEL;TYPE=cell;VALUE=uri:tel:+{}\n'.format(phone_number)) + file.write('END:VCARD\n') + + return file_path + + def download_file_loc(self, input_location, file_path, part_size_kb=64): """Downloads media from the given input_file_location to the specified file_path""" part_size = int(part_size_kb * 1024) if part_size % 1024 != 0: raise ValueError('The part size must be evenly divisible by 1024') + # Ensure that we'll be able to download the media + utils.ensure_parent_dir_exists(file_path) + # Start with an offset index of 0 offset_index = 0 with open(file_path, 'wb') as file: while True: # The current offset equals the offset_index multiplied by the part size offset = offset_index * part_size - result = self.invoke(GetFileRequest(input_file_location, offset, part_size)) + result = self.invoke(GetFileRequest(input_location, offset, part_size)) offset_index += 1 # If we have received no data (0 bytes), the file is over diff --git a/utils/helpers.py b/utils/helpers.py index badccb91..1b49519d 100755 --- a/utils/helpers.py +++ b/utils/helpers.py @@ -1,4 +1,5 @@ import os +import shutil from utils import BinaryWriter import hashlib @@ -30,6 +31,14 @@ def load_settings(path='api/settings'): return settings + +def ensure_parent_dir_exists(file_path): + """Ensures that the parent directory exists""" + parent = os.path.dirname(file_path) + if parent: + os.makedirs(parent, exist_ok=True) + + # endregion # region Cryptographic related utils