Initial commit for RedAirsoft VK to TG and MAX poster
This commit is contained in:
@@ -0,0 +1,274 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional, Set
|
||||
import aiohttp
|
||||
from loguru import logger
|
||||
from .config import settings
|
||||
|
||||
|
||||
# Global registry of paths currently being written / downloaded / processed
|
||||
_ACTIVE_LOCK = asyncio.Lock()
|
||||
_ACTIVE_PATHS: Set[str] = set()
|
||||
|
||||
|
||||
async def register_active_path(path: Path) -> None:
|
||||
async with _ACTIVE_LOCK:
|
||||
_ACTIVE_PATHS.add(str(path.resolve()))
|
||||
|
||||
|
||||
async def unregister_active_path(path: Path) -> None:
|
||||
async with _ACTIVE_LOCK:
|
||||
_ACTIVE_PATHS.discard(str(path.resolve()))
|
||||
|
||||
|
||||
async def is_path_active(path: Path) -> bool:
|
||||
async with _ACTIVE_LOCK:
|
||||
return str(path.resolve()) in _ACTIVE_PATHS
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProcessedMedia:
|
||||
media_type: str # "photo" or "video"
|
||||
local_path: Optional[Path]
|
||||
original_url: str
|
||||
attachment_id: str
|
||||
size_bytes: int = 0
|
||||
duration_sec: Optional[int] = None
|
||||
width: Optional[int] = None
|
||||
height: Optional[int] = None
|
||||
error: Optional[str] = None
|
||||
is_link_only: bool = False
|
||||
|
||||
|
||||
class MediaProcessor:
|
||||
def __init__(self, is_local_tg_api: bool = False) -> None:
|
||||
self.is_local_tg_api = is_local_tg_api
|
||||
self.cache_dir = settings.cache_path
|
||||
|
||||
@property
|
||||
def max_video_bytes(self) -> int:
|
||||
max_mb = (
|
||||
settings.video_max_size_mb_local
|
||||
if self.is_local_tg_api
|
||||
else settings.video_max_size_mb_cloud
|
||||
)
|
||||
return max_mb * 1024 * 1024
|
||||
|
||||
async def download_photo(
|
||||
self, session: aiohttp.ClientSession, url: str, prefix: str = "photo_"
|
||||
) -> Optional[Path]:
|
||||
if not url.startswith(("http://", "https://")):
|
||||
return None
|
||||
|
||||
suffix = Path(url.split("?", 1)[0]).suffix[:6] or ".jpg"
|
||||
tmp = tempfile.NamedTemporaryFile(
|
||||
dir=self.cache_dir, prefix=prefix, suffix=suffix, delete=False
|
||||
)
|
||||
tmp_path = Path(tmp.name)
|
||||
tmp.close()
|
||||
|
||||
await register_active_path(tmp_path)
|
||||
try:
|
||||
async with session.get(
|
||||
url, timeout=aiohttp.ClientTimeout(total=settings.media_download_timeout_sec)
|
||||
) as resp:
|
||||
if resp.status != 200:
|
||||
logger.warning("Photo download failed with status {}: {}", resp.status, url)
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
await unregister_active_path(tmp_path)
|
||||
return None
|
||||
with open(tmp_path, "wb") as f:
|
||||
async for chunk in resp.content.iter_chunked(64 * 1024):
|
||||
f.write(chunk)
|
||||
return tmp_path
|
||||
except Exception as exc:
|
||||
logger.warning("Error downloading photo {}: {}", url, exc)
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
await unregister_active_path(tmp_path)
|
||||
return None
|
||||
|
||||
async def download_video_ytdlp(
|
||||
self, url: str, prefix: str = "video_"
|
||||
) -> tuple[Optional[Path], Optional[str], bool]:
|
||||
"""Returns (path, error_message, is_permanent_error)"""
|
||||
tmp = tempfile.NamedTemporaryFile(
|
||||
dir=self.cache_dir, prefix=prefix, suffix=".mp4", delete=False
|
||||
)
|
||||
output_path = Path(tmp.name)
|
||||
tmp.close()
|
||||
|
||||
await register_active_path(output_path)
|
||||
max_size = self.max_video_bytes
|
||||
netrc_path = None
|
||||
|
||||
cmd = [sys.executable, "-m", "yt_dlp"]
|
||||
if "vk.com" in url and settings.vk_access_token:
|
||||
netrc_path = f"{output_path}.netrc"
|
||||
try:
|
||||
with open(netrc_path, "w", encoding="utf-8") as fh:
|
||||
fh.write(f"machine vk.com login vk_token password {settings.vk_access_token}\n")
|
||||
cmd.extend(["--netrc-location", netrc_path])
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to write netrc for yt-dlp: {}", exc)
|
||||
|
||||
cmd.extend([
|
||||
url,
|
||||
"-o", str(output_path),
|
||||
"--no-playlist",
|
||||
"--match-filter", f"duration <= {settings.video_max_duration_sec}",
|
||||
"--merge-output-format", "mp4",
|
||||
"-f", (
|
||||
f"best[height<={settings.video_max_height}][filesize<{max_size}]"
|
||||
f"/best[height<={settings.video_max_height}]"
|
||||
f"/bestvideo[height<={settings.video_max_height}][filesize<{max_size}]+bestaudio/best"
|
||||
f"/best[filesize<{max_size}]"
|
||||
f"/best"
|
||||
),
|
||||
"--quiet", "--no-warnings",
|
||||
])
|
||||
|
||||
proc = None
|
||||
stderr = b""
|
||||
try:
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
_, stderr = await asyncio.wait_for(
|
||||
proc.communicate(), timeout=settings.yt_dlp_timeout_sec
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
if proc:
|
||||
try:
|
||||
proc.kill()
|
||||
await proc.communicate()
|
||||
except Exception:
|
||||
pass
|
||||
output_path.unlink(missing_ok=True)
|
||||
await unregister_active_path(output_path)
|
||||
return None, "video download timeout", False
|
||||
finally:
|
||||
if netrc_path:
|
||||
Path(netrc_path).unlink(missing_ok=True)
|
||||
|
||||
if proc.returncode != 0 or not output_path.exists() or output_path.stat().st_size == 0:
|
||||
err_msg = (stderr or b"").decode("utf-8", errors="ignore").lower()
|
||||
is_permanent = any(
|
||||
m in err_msg for m in (
|
||||
"removed", "unavailable", "private", "access denied", "does not pass filter", "sign in"
|
||||
)
|
||||
)
|
||||
output_path.unlink(missing_ok=True)
|
||||
await unregister_active_path(output_path)
|
||||
return None, f"yt-dlp failed: {err_msg[:200]}", is_permanent
|
||||
|
||||
size = output_path.stat().st_size
|
||||
if size > max_size:
|
||||
logger.warning(
|
||||
"Downloaded video size {} MB exceeds limit {} MB",
|
||||
round(size / (1024 * 1024), 2),
|
||||
round(max_size / (1024 * 1024), 2),
|
||||
)
|
||||
output_path.unlink(missing_ok=True)
|
||||
await unregister_active_path(output_path)
|
||||
return None, "video exceeds size limit", True
|
||||
|
||||
return output_path, None, False
|
||||
|
||||
async def process_media_items(
|
||||
self, items: list[Any]
|
||||
) -> list[ProcessedMedia]:
|
||||
results: list[ProcessedMedia] = []
|
||||
async with aiohttp.ClientSession() as session:
|
||||
for item in items:
|
||||
media_type = item.media_type
|
||||
url = item.url
|
||||
att_id = item.attachment_id
|
||||
|
||||
if media_type == "photo":
|
||||
p = await self.download_photo(session, url, prefix=f"p_{att_id}_")
|
||||
if p:
|
||||
results.append(
|
||||
ProcessedMedia(
|
||||
media_type="photo",
|
||||
local_path=p,
|
||||
original_url=url,
|
||||
attachment_id=att_id,
|
||||
size_bytes=p.stat().st_size,
|
||||
width=item.width,
|
||||
height=item.height,
|
||||
)
|
||||
)
|
||||
else:
|
||||
results.append(
|
||||
ProcessedMedia(
|
||||
media_type="photo",
|
||||
local_path=None,
|
||||
original_url=url,
|
||||
attachment_id=att_id,
|
||||
error="failed to download photo",
|
||||
is_link_only=True,
|
||||
)
|
||||
)
|
||||
elif media_type == "video":
|
||||
# Check duration
|
||||
if item.duration_sec and item.duration_sec > settings.video_max_duration_sec:
|
||||
results.append(
|
||||
ProcessedMedia(
|
||||
media_type="video",
|
||||
local_path=None,
|
||||
original_url=url,
|
||||
attachment_id=att_id,
|
||||
duration_sec=item.duration_sec,
|
||||
error="video duration exceeds maximum",
|
||||
is_link_only=True,
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
p, err, perm = await self.download_video_ytdlp(url, prefix=f"v_{att_id}_")
|
||||
if p:
|
||||
results.append(
|
||||
ProcessedMedia(
|
||||
media_type="video",
|
||||
local_path=p,
|
||||
original_url=url,
|
||||
attachment_id=att_id,
|
||||
size_bytes=p.stat().st_size,
|
||||
duration_sec=item.duration_sec,
|
||||
width=item.width,
|
||||
height=item.height,
|
||||
)
|
||||
)
|
||||
else:
|
||||
results.append(
|
||||
ProcessedMedia(
|
||||
media_type="video",
|
||||
local_path=None,
|
||||
original_url=url,
|
||||
attachment_id=att_id,
|
||||
duration_sec=item.duration_sec,
|
||||
error=err or "failed to download video",
|
||||
is_link_only=True,
|
||||
)
|
||||
)
|
||||
return results
|
||||
|
||||
async def cleanup(self, items: list[ProcessedMedia]) -> None:
|
||||
"""Immediately unregisters and removes downloaded local files."""
|
||||
for item in items:
|
||||
if item.local_path:
|
||||
try:
|
||||
await unregister_active_path(item.local_path)
|
||||
item.local_path.unlink(missing_ok=True)
|
||||
except Exception as exc:
|
||||
logger.warning("Error cleaning up {}: {}", item.local_path, exc)
|
||||
item.local_path = None
|
||||
Reference in New Issue
Block a user