-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
94 lines (83 loc) · 3.22 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
import os
import uuid
from fastapi import FastAPI, File, Request, UploadFile
from fastapi.responses import StreamingResponse
from fastapi.responses import FileResponse
from starlette.middleware.authentication import AuthenticationMiddleware
import subprocess
import shlex
import logging
from fastapi.middleware.cors import CORSMiddleware
import tempfile
import magic
from middlewares.authentication_middleware import BasicAuthMiddlewareBackend
from middlewares.limit_upload_size import LimitUploadSize
from slowapi.errors import RateLimitExceeded
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
file_upload_limit = int(os.environ.get('FILE_UPLOAD_LIMIT', '50_000_000'))
limiter = Limiter(key_func=get_remote_address)
app = FastAPI()
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
app.add_middleware(
CORSMiddleware,
allow_origins=['*'],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
app.add_middleware(
LimitUploadSize, max_upload_size=file_upload_limit
)
if os.path.exists('.htpasswd'):
app.add_middleware(AuthenticationMiddleware, backend=BasicAuthMiddlewareBackend())
@app.post('/convert')
@limiter.limit("5/minute")
def convert(request: Request, output_options: str, file: UploadFile = File(...), input_options: str = '', buffer_size: int=524288):
try:
command = f"ffmpeg -loglevel error {input_options} -i /dev/stdin {output_options} -"
ffmpeg_cmd = subprocess.Popen(
shlex.split(command),
stdin=file.file,
stdout=subprocess.PIPE,
shell=False
)
mime_type = 'application/octet-stream'
def generate():
while True:
output = ffmpeg_cmd.stdout.read(buffer_size)
if len(output) > 0:
mime_type = magic.Magic(mime=True).from_buffer(output)
yield output
else:
ffmpeg_cmd.wait()
last = ffmpeg_cmd.stdout.read(buffer_size)
yield last
break
return StreamingResponse(generate(), media_type=mime_type)
except Exception as e:
logging.error(e)
if ffmpeg_cmd is not None:
ffmpeg_cmd.terminate()
return {'error': str(e) }
@app.post('/convert-file')
@limiter.limit("5/minute")
def convert(request: Request, output_options: str, file: UploadFile = File(...), input_options: str = ''):
try:
suffix = output_options.split(' ')[output_options.split(' ').index('-f') + 1]
temp_df, temp_path = tempfile.mkstemp(suffix, prefix=str(uuid.uuid4()))
command = f"ffmpeg -loglevel error -y {input_options} -i /dev/stdin {output_options} {temp_path}"
ffmpeg_cmd = subprocess.Popen(
shlex.split(command),
stdin=file.file,
shell=False,
)
ffmpeg_cmd.wait()
mime_type = magic.Magic(mime=True).from_file(temp_path)
return FileResponse(temp_path, media_type=f'{mime_type}')
except Exception as e:
logging.error(e)
if ffmpeg_cmd is not None:
ffmpeg_cmd.terminate()
return {'error': str(e) }