poke/videobundler/main.py

105 lines
3.5 KiB
Python
Raw Normal View History

import asyncio
import aiohttp
from aiohttp import web
import string
2024-06-22 20:52:41 +02:00
import os
import random
import subprocess
2024-06-23 22:11:12 +02:00
from aiohttp.web import Response, FileResponse
2024-06-22 20:52:41 +02:00
app = web.Application()
app.router._frozen = False
2024-06-22 20:52:41 +02:00
def get_random_string(length):
2024-06-24 09:55:25 +02:00
# choose from all lowercase letter
letters = string.ascii_lowercase
result_str = "".join(random.choice(letters) for i in range(length))
return result_str
2024-06-24 05:28:08 +02:00
async def run_command(cmd):
2024-06-24 09:55:25 +02:00
# Create subprocess
process = await asyncio.create_subprocess_shell(
cmd,
stdout=asyncio.subprocess.PIPE,
)
# Wait for the subprocess to finish
stdout, stderr = await process.communicate()
# Check for errors
if process.returncode!= 0:
# Log or handle the error
print(f"Command '{args}' failed with return code {process.returncode}")
return None
# Decode stdout and return
return stdout
2024-06-23 00:02:30 +02:00
2024-06-24 09:55:25 +02:00
async def merge(request: aiohttp.web.Request):
# register params
video_id: str = request.rel_url.query["id"]
audio_itag: str = request.rel_url.query["audio_itag"]
video_itag: str = request.rel_url.query["video_itag"]
# validate
if " " in video_id or len(video_id) > 11:
print(f"Video {video_id} flagged as invalid, dropping request")
return
if not audio_itag.isdigit():
print(f"Audio itag {audio_itag} flagged as invalid, dropping request")
return
if not video_itag.isdigit():
print(f"Video itag {video_itag} flagged as invalid, dropping request")
return
if "Firefox" in request.headers["User-Agent"]:
# Sane browser that supports streaming
cmdline = f"ffmpeg -i \"https://eu-proxy.poketube.fun/latest_version?id={video_id}&itag={audio_itag}&local=true\" -i \"https://eu-proxy.poketube.fun/latest_version?id={video_id}&itag={video_itag}&local=true\" -c copy -f mp4 -movflags frag_keyframe+empty_moov -"
process = await asyncio.create_subprocess_shell(
cmdline,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE
)
response = web.StreamResponse(status=206, reason='OK', headers={
2024-06-25 20:34:41 +02:00
'Content-Type': 'application/octet-stream',
'Transfer-Encoding': 'chunked',
'Content-Disposition': 'inline'
2024-06-24 09:55:25 +02:00
})
await response.prepare(request)
try:
while True:
chunk = await process.stdout.readline()
if not chunk:
break
await response.write(chunk)
except Exception as e:
2024-06-25 20:34:41 +02:00
2024-06-24 09:55:25 +02:00
print(f"Error streaming FFmpeg output: {e}")
2024-06-25 20:34:41 +02:00
#finally:
#await response.write_eof()
2024-06-24 09:55:25 +02:00
else:
# Likely to be chromium browser, so to avoid browser shitting itself we download file
2024-06-26 12:17:04 +02:00
job_id = f'{request.rel_url.query["id"]}_{request.rel_url.query["audio_itag"]}_{request.rel_url.query["video_itag"]}'
2024-06-24 09:55:25 +02:00
if os.path.isfile(f"{job_id}.mp4"):
return web.FileResponse(
path=f"{job_id}.mp4"
)
2024-06-25 20:34:41 +02:00
cmdline = f"ffmpeg -i \"https://eu-proxy.poketube.fun/latest_version?id={video_id}&itag={audio_itag}&local=true\" -i \"https://eu-proxy.poketube.fun/latest_version?id={video_id}&itag={video_itag}&local=true\" -c:v copy -f mp4 -movflags frag_keyframe+empty_moov {job_id}.mp4"
2024-06-24 09:55:25 +02:00
process = await asyncio.create_subprocess_shell(
2024-06-25 20:34:41 +02:00
cmdline
2024-06-24 09:55:25 +02:00
)
await process.wait()
if process.returncode != 0: # Log or handle the error
return None
response = FileResponse(path=f"{job_id}.mp4")
return response
2024-06-22 20:52:41 +02:00
async def ping(request):
2024-06-24 09:55:25 +02:00
return web.Response(body='{"success": true}', content_type="application/json")
2024-06-22 20:52:41 +02:00
async def init_app():
2024-06-24 09:55:25 +02:00
app.router.add_get("/{id:.+}", merge)
app.router.add_get("/", ping)
return app
2024-06-23 00:02:30 +02:00
if __name__ == '__main__':
2024-06-24 09:55:25 +02:00
loop = asyncio.get_event_loop()
app = loop.run_until_complete(init_app())
web.run_app(app, port=3030)