143 lines
4.7 KiB
Python
143 lines
4.7 KiB
Python
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]
|