Reorganize the project for better compatibility with unittests

telegram_downloader_bot/utils.py

@@ -0,0 +1,142 @@
+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]