247 lines
10 KiB
Python
247 lines
10 KiB
Python
"""Stream Telegram media to disk with resume, disk-space checks and hashing."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import errno
|
|
import hashlib
|
|
import logging
|
|
import os
|
|
import re
|
|
import shutil
|
|
from pathlib import Path
|
|
from typing import AsyncIterator, Callable, Protocol
|
|
|
|
from errors import DiskSpaceError, Interrupted, TransientError
|
|
from models import LocalFile, State, VideoRecord
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
# Partial downloads are resumed from a 1 MiB boundary: Telegram's upload.getFile
|
|
# requires offsets aligned to the requested chunk size (512 KiB).
|
|
RESUME_ALIGNMENT = 1024 * 1024
|
|
HASH_BLOCK = 4 * 1024 * 1024
|
|
# Only files with exactly this shape are ever deleted by cleanup routines.
|
|
TEMP_NAME = re.compile(r"^(\d+)_(\d+)\.(part|video)$")
|
|
|
|
|
|
class MediaSource(Protocol):
|
|
def iter_media_chunks(
|
|
self, message_id: int, *, offset: int, expected_size: int | None, channel_id: int | None = None
|
|
) -> AsyncIterator[bytes]: ...
|
|
|
|
|
|
class _Hashes:
|
|
def __init__(self, sha256: bool):
|
|
self.md5 = hashlib.md5()
|
|
self.sha256 = hashlib.sha256() if sha256 else None
|
|
|
|
def update(self, data: bytes) -> None:
|
|
self.md5.update(data)
|
|
if self.sha256:
|
|
self.sha256.update(data)
|
|
|
|
def digest(self) -> tuple[str, str | None]:
|
|
return self.md5.hexdigest(), self.sha256.hexdigest() if self.sha256 else None
|
|
|
|
|
|
def hash_file(path: Path, *, sha256: bool, limit: int | None = None, hashes: _Hashes | None = None) -> _Hashes:
|
|
hashes = hashes or _Hashes(sha256)
|
|
remaining = limit
|
|
with open(path, "rb") as fh:
|
|
while remaining is None or remaining > 0:
|
|
block = fh.read(HASH_BLOCK if remaining is None else min(HASH_BLOCK, remaining))
|
|
if not block:
|
|
break
|
|
hashes.update(block)
|
|
if remaining is not None:
|
|
remaining -= len(block)
|
|
return hashes
|
|
|
|
|
|
class Downloader:
|
|
def __init__(
|
|
self,
|
|
source: MediaSource,
|
|
download_dir: Path,
|
|
*,
|
|
min_free_bytes: int = 0,
|
|
max_temp_bytes: int = 0,
|
|
compute_sha256: bool = True,
|
|
disk_usage: Callable[[str], tuple] = shutil.disk_usage,
|
|
):
|
|
self.source = source
|
|
self.download_dir = download_dir
|
|
self.min_free_bytes = min_free_bytes
|
|
self.max_temp_bytes = max_temp_bytes
|
|
self.compute_sha256 = compute_sha256
|
|
self._disk_usage = disk_usage
|
|
# Optional callback that frees space by removing abandoned temp files;
|
|
# returns the number of files deleted. Set by the queue manager.
|
|
self.reclaim: Callable[[], int] | None = None
|
|
download_dir.mkdir(parents=True, exist_ok=True)
|
|
self._root = download_dir.resolve()
|
|
|
|
# Local files are named by message identity, never by caption text, so
|
|
# captions cannot cause path problems.
|
|
def complete_path(self, rec: VideoRecord) -> Path:
|
|
return self.download_dir / f"{abs(rec.channel_id)}_{rec.message_id}.video"
|
|
|
|
def partial_path(self, rec: VideoRecord) -> Path:
|
|
return self.download_dir / f"{abs(rec.channel_id)}_{rec.message_id}.part"
|
|
|
|
def has_complete_local(self, rec: VideoRecord) -> bool:
|
|
path = self.complete_path(rec)
|
|
if not path.is_file():
|
|
return False
|
|
expected = rec.expected_size or rec.local_size
|
|
return expected is not None and path.stat().st_size == expected
|
|
|
|
async def local_file(self, rec: VideoRecord) -> LocalFile:
|
|
"""Describe an already complete download, reusing stored hashes if valid."""
|
|
path = self.complete_path(rec)
|
|
size = path.stat().st_size
|
|
if rec.md5 and rec.local_size == size:
|
|
return LocalFile(path, size, rec.md5, rec.sha256)
|
|
hashes = await asyncio.to_thread(hash_file, path, sha256=self.compute_sha256)
|
|
md5, sha = hashes.digest()
|
|
return LocalFile(path, size, md5, sha)
|
|
|
|
def check_space(self, needed: int, *, exclude: Path | None = None) -> None:
|
|
free = self._disk_usage(str(self.download_dir))[2]
|
|
if free - needed < self.min_free_bytes:
|
|
raise DiskSpaceError(
|
|
f"not enough disk space: need {needed} bytes plus {self.min_free_bytes} reserve, {free} free",
|
|
public="Not enough free disk space for the next video; free space and /resume.",
|
|
)
|
|
if self.max_temp_bytes:
|
|
used = sum(
|
|
p.stat().st_size for p in self.download_dir.iterdir() if p.is_file() and p != exclude
|
|
)
|
|
if used + needed > self.max_temp_bytes:
|
|
raise DiskSpaceError(
|
|
f"MAX_TEMP_STORAGE exceeded: {used} used + {needed} needed > {self.max_temp_bytes}",
|
|
public="Temporary storage limit (MAX_TEMP_STORAGE) reached; /retry or clean up downloads.",
|
|
)
|
|
|
|
async def download(
|
|
self,
|
|
rec: VideoRecord,
|
|
*,
|
|
on_progress: Callable[[int, int | None], None],
|
|
should_stop: Callable[[], bool],
|
|
) -> LocalFile:
|
|
if self.has_complete_local(rec):
|
|
return await self.local_file(rec)
|
|
part = self.partial_path(rec)
|
|
final = self.complete_path(rec)
|
|
expected = rec.expected_size
|
|
offset = part.stat().st_size if part.exists() else 0
|
|
offset -= offset % RESUME_ALIGNMENT
|
|
if expected is not None and offset > expected:
|
|
offset = 0
|
|
needed = (expected - offset) if expected else 0
|
|
try:
|
|
self.check_space(needed, exclude=part)
|
|
except DiskSpaceError:
|
|
if not (self.reclaim and self.reclaim()):
|
|
raise
|
|
self.check_space(needed, exclude=part)
|
|
|
|
hashes = _Hashes(self.compute_sha256)
|
|
if offset:
|
|
log.info("Resuming download of message %s at byte %d", rec.message_id, offset)
|
|
await asyncio.to_thread(hash_file, part, sha256=self.compute_sha256, limit=offset, hashes=hashes)
|
|
written = offset
|
|
chunks = self.source.iter_media_chunks(
|
|
rec.message_id, offset=offset, expected_size=expected, channel_id=rec.channel_id
|
|
)
|
|
try:
|
|
with open(part, "r+b" if part.exists() else "wb") as fh:
|
|
fh.truncate(offset)
|
|
fh.seek(offset)
|
|
on_progress(written, expected)
|
|
async for chunk in chunks:
|
|
fh.write(chunk)
|
|
hashes.update(chunk)
|
|
written += len(chunk)
|
|
on_progress(written, expected)
|
|
if expected is not None and written > expected:
|
|
break
|
|
if should_stop():
|
|
fh.flush()
|
|
raise Interrupted(f"download paused at byte {written}")
|
|
fh.flush()
|
|
os.fsync(fh.fileno())
|
|
except OSError as exc:
|
|
if exc.errno == errno.ENOSPC:
|
|
raise DiskSpaceError("disk full while downloading", public="Local disk is full; free space and /resume.") from exc
|
|
raise TransientError(f"download failed: {exc}", public="Download failed (I/O or network error)") from exc
|
|
finally:
|
|
await chunks.aclose()
|
|
|
|
size = part.stat().st_size
|
|
if expected is not None and size != expected:
|
|
if size > expected:
|
|
part.unlink(missing_ok=True)
|
|
raise TransientError(
|
|
f"incomplete download: {size} of {expected} bytes", public="Download incomplete or corrupted; will retry"
|
|
)
|
|
if size == 0:
|
|
part.unlink(missing_ok=True)
|
|
raise TransientError("downloaded file is empty", public="Downloaded file is empty; will retry")
|
|
os.replace(part, final)
|
|
md5, sha = hashes.digest()
|
|
return LocalFile(final, size, md5, sha)
|
|
|
|
def cleanup(self, rec: VideoRecord) -> int:
|
|
"""Delete this video's temporary files. Returns how many files were removed."""
|
|
return sum(self.safe_unlink(path) for path in (self.complete_path(rec), self.partial_path(rec)))
|
|
|
|
def safe_unlink(self, path: Path) -> bool:
|
|
"""Delete ``path`` only if it is a regular temp file directly inside the download dir."""
|
|
try:
|
|
if path.is_symlink() or path.parent.resolve() != self._root or not TEMP_NAME.match(path.name):
|
|
log.error("Refusing to delete %s: not an application temp file", path.name)
|
|
return False
|
|
if not path.exists():
|
|
return False
|
|
path.unlink()
|
|
log.info("Deleted temporary file %s", path.name)
|
|
return True
|
|
except OSError as exc:
|
|
log.warning("Could not delete temporary file %s: %s", path.name, exc)
|
|
return False
|
|
|
|
def cleanup_orphans(self, db, *, keep_failed: bool = False) -> int:
|
|
"""Remove temp files that no unfinished item needs any more.
|
|
|
|
A file is deleted only when it matches the application's temp-name pattern,
|
|
lives directly in the download directory (symlinks are never followed), and
|
|
its database record is missing or finished (COMPLETED / SKIPPED, or FAILED
|
|
unless ``keep_failed``). Leftover ``.part`` files of fully downloaded items
|
|
are removed too. Anything else in the directory is left untouched.
|
|
"""
|
|
deleted = 0
|
|
try:
|
|
entries = list(os.scandir(self._root))
|
|
except FileNotFoundError:
|
|
return 0
|
|
for entry in entries:
|
|
match = TEMP_NAME.match(entry.name)
|
|
if not match or entry.is_symlink() or not entry.is_file(follow_symlinks=False):
|
|
continue
|
|
channel_id, message_id, kind = -int(match.group(1)), int(match.group(2)), match.group(3)
|
|
rec = db.get((channel_id, message_id))
|
|
if rec is None or rec.state in (State.COMPLETED, State.SKIPPED):
|
|
reason = "record finished" if rec else "no database record"
|
|
elif rec.state == State.FAILED and not keep_failed:
|
|
reason = "item failed (can be downloaded again)"
|
|
elif kind == "part" and rec.state in (State.DOWNLOADED, State.UPLOADING):
|
|
reason = "stale partial of a completed download"
|
|
else:
|
|
continue
|
|
if self.safe_unlink(Path(entry.path)):
|
|
deleted += 1
|
|
log.info("Orphan cleanup: removed %s (%s)", entry.name, reason)
|
|
return deleted
|