import integv import os import re from datetime import datetime from pyrogram import Client from pyrogram.types import Message from tiktok_downloader import snaptik from telegram_downloader_bot.logger import log async def get_user_folder(storage_path: os.path, message: Message) -> os.path: """ Determine folder name used to save the media to. Depending on which type of message (forwarded, direct) detect that person's or group's name.""" if message.forward_from: user = message.forward_from # User's first and last name for folder name, # fallback to user ID if not available user_folder_name = ( f"{user.first_name}_{user.last_name}".strip() if user.first_name and user.last_name else str(user.id) ) elif message.forward_from_chat: user = message.forward_from_chat # Use chat title for groups and channels user_folder_name = "".join( c for c in user.title if c.isalnum() or c in (" ", "_") ).rstrip() else: user = message.from_user # User's first and last name for folder name, # fallback to user ID if not available user_folder_name = ( f"{user.first_name}_{user.last_name}".strip() if user.first_name and user.last_name else str(user.id) ) # Sanitize the folder name user_folder_name = "".join( c for c in user_folder_name if c.isalnum() or c in (" ", "_") ).rstrip() user_folder = os.path.join(storage_path, "telegram", user_folder_name) os.makedirs(user_folder, exist_ok=True) return user_folder async def handle_media_message_contents(storage_path: os.path, client: Client, message: Message): """Detect what kind of media is being sent over from the user. Based on that, determine the correct file extension and save that media.""" user_folder = get_user_folder(storage_path, message) # Handle documents if message.document: file_name = message.document.file_name file_path = os.path.join(user_folder, file_name) await client.download_media(message, file_path) await message.reply_text(f"Document saved to {user_folder}") # Handle single or multiple photos elif message.photo: file_name = f"photo_{message.photo.file_id}.jpg" file_path = os.path.join(user_folder, file_name) await client.download_media(message.photo, file_path) await message.reply_text(f"Photo saved to {user_folder}") # Handle videos elif message.video: file_name = f"video_{message.video.file_id}.mp4" file_path = os.path.join(user_folder, file_name) await client.download_media(message, file_path) await message.reply_text(f"Video saved to {user_folder}") # Handle GIFs elif message.animation: file_name = f"gif_{message.animation.file_id}.gif" file_path = os.path.join(user_folder, file_name) await client.download_media(message.animation, file_path) await message.reply_text(f"GIF saved to {user_folder}") # Handle unknown data types else: await message.reply_text("Unknown media type!") def download_tt_video(storage_path: str, url: str) -> bool: """Downloads tiktok video from a given URL. Makes sure the video integrity is correct.""" videos = snaptik(url) now = datetime.datetime.now() max_tries = 5 log.debug(f"Downloading video from {url}...") for video in videos: video_filename = now.strftime("video-tiktok-%Y-%m-%d_%H-%M-%S.mp4") video_filepath: os.path = os.path.join( storage_path, "tiktok", video_filename) for i in range(max_tries): video_content = video.download().getbuffer() is_valid_mp4 = integv.verify( video_content.tobytes(), file_type="mp4") log.debug( f"Attempt {i+1}/{max_tries} to " "download video, video valid: {is_valid_mp4}" ) if is_valid_mp4: break if not is_valid_mp4: log.error("Downloaded video is not a valid mp4 file") return False with open(video_filepath, "wb") as f: f.write(video_content) log.debug("Video saved successfully") return True def make_fs(storaga_path: str) -> None: os.makedirs(os.path.join(storaga_path, "tiktok"), exist_ok=True) os.makedirs(os.path.join(storaga_path, "telegram"), exist_ok=True) def extract_urls(text: str) -> list: return re.findall(r"\bhttps?://[^\s]+", text) def filter_tt_urls(urls: list) -> list: return [x for x in urls if "tiktok" in x]