198 lines
7.5 KiB
Python
198 lines
7.5 KiB
Python
import os
|
||
import uuid
|
||
import subprocess
|
||
import requests
|
||
import tempfile
|
||
import shutil
|
||
import logging
|
||
from contextlib import contextmanager
|
||
import time
|
||
from shared_utils import add_task_log
|
||
|
||
@contextmanager
|
||
def temp_directory(dir=None):
|
||
temp_dir = tempfile.mkdtemp(dir=dir)
|
||
try:
|
||
yield temp_dir
|
||
finally:
|
||
pass
|
||
|
||
def download_file(url, temp_dir, task_id=None, task_registry=None, task_lock=None):
|
||
try:
|
||
headers = {
|
||
'User-Agent': 'Mozilla/5.0',
|
||
'Accept': '*/*'
|
||
}
|
||
|
||
if task_id and task_registry and task_lock:
|
||
add_task_log(task_id, f"开始从URL下载文件: {url}", task_registry, task_lock)
|
||
|
||
logging.info(f"开始从URL下载文件: {url}")
|
||
response = requests.get(url, headers=headers, stream=True, timeout=30)
|
||
response.raise_for_status()
|
||
|
||
file_path = os.path.join(temp_dir, 'input_audio')
|
||
|
||
with open(file_path, 'wb') as f:
|
||
for chunk in response.iter_content(chunk_size=8192):
|
||
if chunk:
|
||
f.write(chunk)
|
||
|
||
file_size = os.path.getsize(file_path)
|
||
if file_size < 1024:
|
||
error_msg = f"下载的文件太小(仅{file_size}字节)，可能无效"
|
||
logging.error(error_msg)
|
||
raise Exception(error_msg)
|
||
|
||
if task_id and task_registry and task_lock:
|
||
add_task_log(task_id, f"文件下载成功，保存到: {file_path} (大小: {file_size}字节)", task_registry, task_lock)
|
||
|
||
logging.info(f"文件下载成功，保存到: {file_path} (大小: {file_size}字节)")
|
||
return file_path
|
||
except Exception as e:
|
||
error_msg = f"从 {url} 下载文件时出错: {e}"
|
||
if task_id and task_registry and task_lock:
|
||
add_task_log(task_id, error_msg, task_registry, task_lock)
|
||
logging.error(error_msg)
|
||
raise
|
||
|
||
def execute_sanjuuni(input_path, output_path, sanjuuni_args, task_id=None, task_registry=None, task_lock=None):
|
||
try:
|
||
cmd = ['lib/sanjuuni/sanjuuni', '-i', input_path] + sanjuuni_args + ['-o', output_path]
|
||
|
||
if task_id and task_registry and task_lock:
|
||
add_task_log(task_id, f"执行Sanjuuni命令: {' '.join(cmd)}", task_registry, task_lock)
|
||
|
||
logging.info(f"执行Sanjuuni命令: {' '.join(cmd)}")
|
||
|
||
# 使用Popen来实时捕获输出
|
||
process = subprocess.Popen(
|
||
cmd,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
universal_newlines=True,
|
||
encoding='utf-8',
|
||
errors='replace',
|
||
bufsize=1
|
||
)
|
||
|
||
# 实时读取输出
|
||
while True:
|
||
output = process.stdout.readline()
|
||
if output == '' and process.poll() is not None:
|
||
break
|
||
if output:
|
||
output = output.strip()
|
||
if output and task_id and task_registry and task_lock:
|
||
add_task_log(task_id, output, task_registry, task_lock)
|
||
|
||
returncode = process.poll()
|
||
if returncode != 0:
|
||
error_msg = f"Sanjuuni处理失败，返回码: {returncode}"
|
||
if task_id and task_registry and task_lock:
|
||
add_task_log(task_id, error_msg, task_registry, task_lock)
|
||
logging.error(error_msg)
|
||
raise Exception(error_msg)
|
||
|
||
if task_id and task_registry and task_lock:
|
||
add_task_log(task_id, "Sanjuuni处理成功完成", task_registry, task_lock)
|
||
|
||
logging.info("Sanjuuni处理成功完成")
|
||
return True
|
||
except Exception as e:
|
||
error_msg = f"执行Sanjuuni时出错: {e}"
|
||
if task_id and task_registry and task_lock:
|
||
add_task_log(task_id, error_msg, task_registry, task_lock)
|
||
logging.error(error_msg)
|
||
raise
|
||
|
||
def process_sanjuuni(data, file_registry, file_lock, task_id=None, task_registry=None, task_lock=None):
|
||
try:
|
||
temp_dir = tempfile.mkdtemp(dir='temp_files')
|
||
|
||
if task_id and task_registry and task_lock:
|
||
add_task_log(task_id, f"创建临时目录: {temp_dir}", task_registry, task_lock)
|
||
|
||
logging.info(f"创建临时目录: {temp_dir}")
|
||
|
||
input_url = data.get('input_url')
|
||
sanjuuni_args = data.get('args', [])
|
||
|
||
# 定义不允许的参数
|
||
disallowed_params = [
|
||
'-s', '--http=', '-w', '--websocket=', '-u', '--websocket-client=',
|
||
'-T', '--streamed', '--disable-opencl', '-i', '--input=', '-o', '--output=',
|
||
'-S', '--subtitles='
|
||
]
|
||
|
||
# 检查是否有不允许的参数
|
||
for arg in sanjuuni_args:
|
||
for param in disallowed_params:
|
||
if arg.startswith(param):
|
||
error_msg = f"不允许使用参数: {arg}"
|
||
logging.error(error_msg)
|
||
raise ValueError(error_msg)
|
||
|
||
subtitle = data.get('subtitle', '')
|
||
output_format = data.get('output_format', 'lua')
|
||
|
||
# 构建 sanjuuni 参数
|
||
if 'format' in data:
|
||
sanjuuni_args.extend(['-f', data['format']])
|
||
|
||
if output_format == 'lua':
|
||
sanjuuni_args.append('-l')
|
||
elif output_format == 'nfp':
|
||
sanjuuni_args.append('-n')
|
||
elif output_format == 'raw':
|
||
sanjuuni_args.append('-r')
|
||
elif output_format == 'bimg':
|
||
sanjuuni_args.append('-b')
|
||
elif output_format == '32vid':
|
||
sanjuuni_args.append('-3')
|
||
else:
|
||
raise ValueError(f"Unsupported output format: {output_format}")
|
||
|
||
input_path = download_file(input_url, temp_dir, task_id, task_registry, task_lock)
|
||
|
||
output_id = str(uuid.uuid4())[:8]
|
||
output_filename = f"{output_id}.{output_format}"
|
||
output_path = os.path.join(temp_dir, output_filename)
|
||
|
||
if task_id and task_registry and task_lock:
|
||
add_task_log(task_id, f"开始Sanjuuni处理，输出格式: {output_format}", task_registry, task_lock)
|
||
with task_lock:
|
||
if task_id in task_registry:
|
||
task_registry[task_id]['progress'] = 50
|
||
|
||
execute_sanjuuni(input_path, output_path, sanjuuni_args, task_id, task_registry, task_lock)
|
||
|
||
with file_lock:
|
||
file_registry[output_id] = {
|
||
'path': os.path.abspath(output_path),
|
||
'filename': output_filename,
|
||
'last_access': time.time(),
|
||
'download_count': 0
|
||
}
|
||
|
||
if task_id and task_registry and task_lock:
|
||
add_task_log(task_id, f"已注册新文件ID: {output_id}, 路径: {output_path}", task_registry, task_lock)
|
||
with task_lock:
|
||
if task_id in task_registry:
|
||
task_registry[task_id]['progress'] = 100
|
||
|
||
logging.info(f"已注册新文件ID: {output_id}, 路径: {output_path}")
|
||
|
||
return {
|
||
'status': 'success',
|
||
'download_url': f"http://newgmapi.liulikeji.cn/download/{output_id}/{output_filename}",
|
||
'file_id': output_id,
|
||
'temp_dir': temp_dir
|
||
}
|
||
|
||
except Exception as e:
|
||
error_msg = f"处理过程中出错: {e}"
|
||
if task_id and task_registry and task_lock:
|
||
add_task_log(task_id, error_msg, task_registry, task_lock)
|
||
logging.error(error_msg)
|
||
return {'error': str(e)} |