TestTG / main.py
rkihacker's picture
Update main.py
69bfb34 verified
raw
history blame
4.25 kB
import os
import asyncio
from telethon.sync import TelegramClient
from telethon.sessions import StringSession
from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
from fastapi.responses import FileResponse
import uvicorn
import tempfile
# --- Configuration from environment variables ---
SESSION_STRING = os.environ.get('SESSION_STRING')
API_ID = os.environ.get('API_ID')
API_HASH = os.environ.get('API_HASH')
# The Telegram channel to use as file storage
TARGET_CHANNEL = 'https://t.me/TestNAI01'
# Use the system's temporary directory for transient files
UPLOAD_DIR = tempfile.gettempdir()
app = FastAPI()
async def upload_to_telegram(file_path: str):
"""Uploads a file and returns the full Telegram message object."""
if not all([SESSION_STRING, API_ID, API_HASH]):
raise ValueError("Server is not configured with Telegram credentials.")
async with TelegramClient(StringSession(SESSION_STRING), API_ID, API_HASH) as client:
try:
message = await client.send_file(TARGET_CHANNEL, file_path, force_document=True)
if not message or not message.media:
raise HTTPException(status_code=500, detail="Upload failed to return a valid message.")
return message
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to upload to Telegram: {e}")
@app.post("/v1/upload")
async def upload_file(file: UploadFile = File(...)):
"""
Uploads a file to Telegram and returns its persistent message_id and other details.
"""
file_location = os.path.join(UPLOAD_DIR, file.filename)
try:
with open(file_location, "wb+") as file_object:
file_object.write(file.file.read())
message = await upload_to_telegram(file_location)
doc = message.document
original_filename = next((attr.file_name for attr in doc.attributes if hasattr(attr, 'file_name')), file.filename)
# Return details without the channel link, as requested.
return {
"message": "File uploaded successfully. Use the 'message_id' to download.",
"telegram_details": {
"message_id": message.id,
"file_id": doc.id,
"file_reference": doc.file_reference.hex() if doc.file_reference else None,
"filename": original_filename,
"size_bytes": doc.size,
"mime_type": doc.mime_type,
"date": message.date.isoformat()
}
}
finally:
if os.path.exists(file_location):
os.remove(file_location)
@app.get("/v1/download/{message_id}")
async def download_file(message_id: int, background_tasks: BackgroundTasks):
"""
Downloads a file from Telegram using its unique message_id from the channel.
This method is stateless and resilient to server restarts.
"""
if not all([SESSION_STRING, API_ID, API_HASH]):
raise HTTPException(status_code=500, detail="Server is not configured for downloads.")
try:
async with TelegramClient(StringSession(SESSION_STRING), API_ID, API_HASH) as client:
message = await client.get_messages(TARGET_CHANNEL, ids=message_id)
if not message or not message.media:
raise HTTPException(status_code=440, detail=f"A message with ID '{message_id}' was not found or contains no media.")
original_filename = message.file.name if message.file and message.file.name else f"download_{message_id}"
download_path = await client.download_media(message, file=UPLOAD_DIR)
# **THE FIX IS HERE:**
# The os.remove is now a background task. It will only run AFTER the
# FileResponse has been fully sent to the client, preventing the error.
if download_path:
background_tasks.add_task(os.remove, download_path)
return FileResponse(download_path, media_type='application/octet-stream', filename=original_filename)
except Exception as e:
# Catch potential errors like file not found in Telegram
raise HTTPException(status_code=500, detail=f"Failed to download from Telegram: {e}")