diff --git a/main.py b/main.py index 5c300ef..9a74203 100644 --- a/main.py +++ b/main.py @@ -8,6 +8,11 @@ from multiprocessing import Pool import shutil +from datetime import datetime +from time import sleep +import yt_dlp +import traceback + vkHandler = None def init(): @@ -53,14 +58,29 @@ def handle_command(obj, command): pic.save(obj) def handler(dir): - imgs = [file for file in os.listdir(dir) if os.path.isfile(dir+"/"+file)] - arch_path = dir+"/arch.zip" - print("handling photo req | dir: ", dir) - with zipfile.ZipFile(arch_path, "w") as archive: - for img in imgs: - handle_command(dir+"/"+img, command) - archive.write(dir+"/"+img, "imgs/"+img) - return arch_path + imgs = sum(list(map(lambda t: ['{}/{}'.format(t,file) for file in os.listdir(dir+"/"+t) if os.path.isfile('{}/{}/{}'.format(dir,t,file))], [directory for directory in os.listdir(dir) if os.path.isdir(dir+"/"+directory)])), []) + print("handling img request | dir: ", dir) + + archives = [] + max_size = 200*1024*1024 + cur_size = 0 + cur_arch = 0 + archive = zipfile.ZipFile(f'{dir}/arch{cur_arch}.zip', 'w') + archives.append(f'{dir}/arch{cur_arch}.zip') + for img in imgs: + img_size = os.path.getsize(f'{dir}/{img}') + if (cur_size + img_size) >= max_size: + cur_size = 0 + cur_arch += 1 + archive.close() + archive = zipfile.ZipFile(f'{dir}/arch{cur_arch}.zip', 'w') + archives.append(f'{dir}/arch{cur_arch}.zip') + #handle_command(f'{dir}/{img}', command) + archive.write(f'{dir}/{img}', img) + cur_size += img_size + archive.close() + + return archives return handler @@ -68,54 +88,74 @@ def img_names_generator(): for name in sequence_generator(): yield str(name)+".jpg" -def handle_request(event, links, request_handler, *args): - request_id = str(generate_id()) - request_dir = "userData/" + request_id - os.mkdir(request_dir) - - vkHandler.send_message('downloading attachments...') - download_files(links, request_dir, *args) - vkHandler.send_message('processing files...') - arch_path = request_handler(request_dir) - vkHandler.send_message('uploading archive...') - vkHandler.send_message(attachment = arch_path) - shutil.rmtree(request_dir) - def download_file(kwargs): - resp = requests.get(kwargs['url']) - with open(kwargs['path'] + "/" + str(kwargs['name']), "wb") as writer: - writer.write(resp.content) - - -def download_files(links, path = ".", name_generator = sequence_generator()): - pool = Pool(max(min(len(links), os.cpu_count()), 1)) - links = [{'url': link, 'path': path, 'name': next(name_generator)} for link in links] - print(links) - pool.map(download_file, links) + if kwargs['type'] in ['photo', 'doc']: + resp = requests.get(kwargs['url']) + if not os.path.isdir(kwargs['path']): + os.mkdir(kwargs['path']) + if kwargs['name'].endswith('.'): + kwargs['name'] += 'jpg' + with open(kwargs['path'] + "/" + str(kwargs['name']), "wb") as writer: + writer.write(resp.content) + + elif kwargs['type'] == 'video' and kwargs['url'] is not None: + if kwargs['name'].endswith('.'): + kwargs['name'] += 'mp4' + ydl_opts = {'outtmpl': kwargs['path'] + "/" + str(kwargs['name']), 'quiet': True} + with yt_dlp.YoutubeDL(ydl_opts) as ydl: + ydl.download([kwargs['url']]) + +def download_files(attachments, path = ".", name_generator = sequence_generator()): + pool = Pool(max(min(len(attachments), os.cpu_count() - 1), 1)) + attachments = [{ + 'type': attachment['type'], + 'url': attachment['url'] if 'url' in attachment else ('https://vk.com/video{}_{}'.format(attachment['owner_id'], attachment['id']) if attachment['type'] == 'video' else None), + 'path': '{}/{}'.format(path,attachment['owner_id']), + 'name': attachment['title'] if "title" in attachment and "ext" in attachment and attachment['title'] is not None and attachment['ext'] is not None and attachment['title'].lower().endswith("."+attachment['ext'].lower()) else ('{}.{}'.format(attachment['title'],attachment['ext']) if "title" in attachment and "ext" in attachment and attachment['title'] is not None and attachment['ext'] is not None else ('{}.'.format(attachment['title']) if 'title' in attachment and attachment['title'] is not None else '{:%Y-%m-%d_%H:%M:%S}_{}.'.format(datetime.fromtimestamp(attachment['date']), attachment['id']))) + } for attachment in attachments] + pool.map(download_file, attachments) pool.close() pool.join() - def bot_loop(): print("Bot loop started.") for event in vkHandler.listen(): if event.type == vkHandler.longpoll.VkBotEventType.MESSAGE_NEW: - photos, command = vkHandler.get_photos_links(), vkHandler.get_command() + attachments, command, message_id = vkHandler.get_attachments(), vkHandler.get_command(), vkHandler.get_message_id() + request_id = str(generate_id()) + request_dir = "userData/" + request_id + os.mkdir(request_dir) + try: - handle_request(event, photos, get_img_request_handler(command), img_names_generator()) - except BaseException as e: - vkHandler.send_message("An error occured while processing your request.\n"+str(e)) + request_handler = get_img_request_handler(command) + vkHandler.send_message('downloading attachments...', reply_to = message_id) + download_files(attachments, request_dir, img_names_generator()) + vkHandler.send_message('processing files...', reply_to = message_id) + archives = request_handler(request_dir) + vkHandler.send_message('uploading archives...', reply_to = message_id) + vkHandler.send_message(attachments = archives, reply_to = message_id) + except BaseException as e: + tb_format = traceback.format_exc() + vkHandler.send_message( + "An error occurred while processing your request.\n{}".format(tb_format.splitlines()[-1]), + reply_to = message_id + ) + print(tb_format) + finally: + shutil.rmtree(request_dir) def main(): init() - print("Bot successfuly inited. Starting bot loop.") + print("Bot successfully initialized. Starting bot loop.") while True: try: + sleep(1/10) bot_loop() except BaseException as e: print(e) - print("An error occured while bot worked. Restarting...") + print("An error occurred during events processing. Restarting...") + sleep(1/10) if __name__ == '__main__': main() diff --git a/vk.py b/vk.py index c19f47b..f9d300d 100644 --- a/vk.py +++ b/vk.py @@ -2,6 +2,8 @@ import vk_api.bot_longpoll as longpoll from random import randint +from datetime import datetime + class vk_wrapper: def __init__(self, access_token, group_id, api_version = '5.103'): @@ -34,10 +36,10 @@ def get_full_message(event, vk_main): return vk_main.method('messages.getById', params)['items'][0] @staticmethod - def get_attachments(message): + def get_message_attachments(message): result = message['attachments'] for fwd_message in message.get('fwd_messages', []): - result.extend(vk_wrapper.get_attachments(fwd_message)) + result.extend(vk_wrapper.get_message_attachments(fwd_message)) return result @@ -45,30 +47,92 @@ def get_attachments(message): def filter_attachments(attachments, at_type): result = [] for attachment in attachments: - if (attachment['type'] == at_type): - result.append(attachment[at_type]) + if (attachment['type'] in at_type): + attachment[attachment['type']]['type'] = attachment['type'] + result.append(attachment[attachment['type']]) return result @staticmethod - def get_photo_link(photo): - sizes = ['s','m','x','o','p','q','r','y','z','w'] - best_size = max(photo['sizes'], key = lambda size: sizes.index(size['type'])) - return best_size['url'] - - def get_photos_links(self): + def get_attachment(attachment): + if attachment['type'] == 'photo': + if attachment['orig_photo'] is not None and attachment['orig_photo']['url'] is not None: + return { + 'id': attachment['id'], + 'owner_id': attachment['owner_id'], + 'date': attachment['date'], + 'url': attachment['orig_photo']['url'], + 'type': attachment['type'] + } + + sizes = ['s','m','x','o','p','q','r','y','z','w'] + best_size = max(attachment['sizes'], key = lambda size: sizes.index(size['type'])) + return { + 'id': attachment['id'], + 'owner_id': attachment['owner_id'], + 'date': attachment['date'], + 'url': best_size['url'], + 'type': attachment['type'] + } + + elif attachment['type'] == 'doc': + if attachment['url'] is not None: + return { + 'id': attachment['id'], + 'owner_id': attachment['owner_id'], + 'date': attachment['date'], + 'url': attachment['url'], + 'title': attachment['title'], + 'ext': attachment['ext'], + 'type': attachment['type'] + } + + sizes = ['s','m','x','y','z','o'] + best_size = max(attachment['preview']['photo']['sizes'], key = lambda size: sizes.index(size['type'])) + return { + 'id': attachment['id'], + 'owner_id': attachment['owner_id'], + 'date': attachment['date'], + 'url': best_size['src'], + 'title': attachment['title'], + 'ext': attachment['ext'], + 'type': attachment['type'] + } + + elif attachment['type'] == 'video': + return { + 'id': attachment['id'], + 'owner_id': attachment['owner_id'], + 'date': attachment['date'], + 'title': attachment['title'] if 'title' in attachment else None, + 'type': attachment['type'] + } + + def get_attachments(self): message = self.get_full_message(self.__event['message'], self.__vk_main) - attachments = self.get_attachments(message) - return [self.get_photo_link(photo) for photo in self.filter_attachments(attachments, 'photo')] + attachments = self.get_message_attachments(message) + return [self.get_attachment(attachment) for attachment in self.filter_attachments(attachments, ['photo', 'doc', 'video'])] def get_command(self): return self.__event['message']['text'] - def get_dialog_id(self): + def get_message_id(self): + return self.__event['message']['id'] + + def get_peer_id(self): return self.__event['message']['peer_id'] - def send_message(self, message = '', attachment = None, ): - peer_id = self.get_dialog_id() + def get_from_id(self): + return self.__event['message']['from_id'] + + def get_id(self): + return self.__event['message']['id'] + + def get_date(self): + return self.__event['message']['date'] + + def send_message(self, message = '', attachments = None, reply_to = None): + peer_id = self.get_peer_id() # try: # if not attachment is None: @@ -77,18 +141,43 @@ def send_message(self, message = '', attachment = None, ): # message += "\nError: can't upload attachment." # attachment = None - if not attachment is None: - attachment = self.__uploader.document_message(attachment, peer_id = peer_id) - - - print('attachment', attachment) - params = { - 'user_id': peer_id, - 'attachment': None if attachment is None else 'doc{owner_id}_{id}'.format(**attachment['doc']), - 'message': message, - 'random_id': randint(0, 99999999999) - } - self.__vk_main.method('messages.send', params) + i = 0 + if not attachments is None and len(attachments) > 0: + for attachment in attachments: + i += 1 + attachment_title = '{}_{}_{:%Y-%m-%d_%H:%M:%S}_{}.zip'.format( + self.get_from_id(), + self.get_id(), + datetime.fromtimestamp(self.get_date()), + i + ) if len(attachments) > 1 else '{}_{}_{:%Y-%m-%d_%H:%M:%S}.zip'.format( + self.get_from_id(), + self.get_id(), + datetime.fromtimestamp(self.get_date()), + ) + attachment = self.__uploader.document_message( + attachment, + peer_id = peer_id, + title = attachment_title + ) + print('attachment', attachment) + + params = { + 'user_id': peer_id, + 'attachment': 'doc{owner_id}_{id}'.format(**attachment['doc']), + 'message': message, + 'random_id': randint(0, 99999999999), + 'reply_to': reply_to + } + self.__vk_main.method('messages.send', params) + + else: + params = { + 'user_id': peer_id, + 'message': message, + 'random_id': randint(0, 99999999999), + 'reply_to': reply_to + } + self.__vk_main.method('messages.send', params) - \ No newline at end of file