257 lines
7.8 KiB
Python
257 lines
7.8 KiB
Python
#!/usr/bin/env python3
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import os
|
||
import socket
|
||
import sys
|
||
import tempfile
|
||
import urllib.error
|
||
import urllib.parse
|
||
import urllib.request
|
||
from pathlib import Path
|
||
from typing import Any, NoReturn, Optional
|
||
|
||
from _docx_common import emit, failure_message, output_file, publish_file
|
||
|
||
|
||
DEFAULT_TIMEOUT_SECONDS = 60
|
||
MAX_ATTACHMENT_BYTES = 25 * 1024 * 1024
|
||
CHUNK_SIZE = 1024 * 1024
|
||
USER_AGENT = "wechat-robot-docx-attachment-downloader/1.0"
|
||
|
||
|
||
class SkillArgumentParser(argparse.ArgumentParser):
|
||
def error(self, message: str) -> NoReturn:
|
||
raise ValueError(f"参数错误：{message}")
|
||
|
||
|
||
def _validate_https_url(value: str) -> str:
|
||
url = value.strip()
|
||
if not url:
|
||
raise ValueError("附件 URL 不能为空")
|
||
if any(character.isspace() or ord(character) < 32 for character in url):
|
||
raise ValueError("附件 URL 不能包含空白字符或控制字符")
|
||
|
||
parsed = urllib.parse.urlsplit(url)
|
||
if parsed.scheme.lower() != "https" or not parsed.hostname:
|
||
raise ValueError("附件 URL 必须是有效的 HTTPS 地址")
|
||
if parsed.username is not None or parsed.password is not None:
|
||
raise ValueError("附件 URL 不允许包含用户名或密码")
|
||
try:
|
||
parsed.port
|
||
except ValueError as exc:
|
||
raise ValueError("附件 URL 端口格式不正确") from exc
|
||
return url
|
||
|
||
|
||
class HTTPSOnlyRedirectHandler(urllib.request.HTTPRedirectHandler):
|
||
def redirect_request(self, req, fp, code, msg, headers, newurl):
|
||
return super().redirect_request(
|
||
req,
|
||
fp,
|
||
code,
|
||
msg,
|
||
headers,
|
||
_validate_https_url(newurl),
|
||
)
|
||
|
||
|
||
def _parse_args(argv: list[str]) -> argparse.Namespace:
|
||
parser = SkillArgumentParser(
|
||
description="下载不超过 25 MiB 的远程 HTTPS 通用附件"
|
||
)
|
||
parser.add_argument(
|
||
"--url",
|
||
"--attachment-url",
|
||
"--attachment_url",
|
||
dest="url",
|
||
required=True,
|
||
help="远程 HTTPS 附件地址",
|
||
)
|
||
parser.add_argument(
|
||
"--output",
|
||
required=True,
|
||
help="附件本地输出路径；允许图片、音视频、压缩包及其他文件类型",
|
||
)
|
||
parser.add_argument(
|
||
"--timeout",
|
||
type=int,
|
||
default=DEFAULT_TIMEOUT_SECONDS,
|
||
help=f"连接和读取超时秒数，默认 {DEFAULT_TIMEOUT_SECONDS}",
|
||
)
|
||
parser.add_argument(
|
||
"--overwrite",
|
||
action="store_true",
|
||
help="允许覆盖本次任务已存在的缓存文件",
|
||
)
|
||
args = parser.parse_args(argv)
|
||
|
||
args.url = _validate_https_url(args.url)
|
||
if args.timeout < 1 or args.timeout > 600:
|
||
raise ValueError("timeout 必须在 1 到 600 秒之间")
|
||
args.output = output_file(args.output, overwrite=args.overwrite)
|
||
return args
|
||
|
||
|
||
def _content_length(headers: Any) -> Optional[int]:
|
||
raw_value = headers.get("Content-Length")
|
||
if raw_value is None:
|
||
return None
|
||
try:
|
||
size = int(raw_value)
|
||
except (TypeError, ValueError):
|
||
return None
|
||
return size if size >= 0 else None
|
||
|
||
|
||
def _reject_if_too_large(size_bytes: int, *, source: str) -> None:
|
||
if size_bytes <= MAX_ATTACHMENT_BYTES:
|
||
return
|
||
raise ValueError(
|
||
"附件超过 25 MiB（26214400 字节）限制，已拒绝下载；"
|
||
f"{source}大小为 {size_bytes} 字节"
|
||
)
|
||
|
||
|
||
def _probe_size(
|
||
opener: urllib.request.OpenerDirector,
|
||
url: str,
|
||
timeout: int,
|
||
) -> Optional[int]:
|
||
request = urllib.request.Request(
|
||
url,
|
||
headers={
|
||
"Accept": "*/*",
|
||
"Accept-Encoding": "identity",
|
||
"User-Agent": USER_AGENT,
|
||
},
|
||
method="HEAD",
|
||
)
|
||
try:
|
||
with opener.open(request, timeout=timeout) as response:
|
||
_validate_https_url(response.geturl())
|
||
size = _content_length(response.headers)
|
||
except (
|
||
urllib.error.HTTPError,
|
||
urllib.error.URLError,
|
||
TimeoutError,
|
||
socket.timeout,
|
||
):
|
||
return None
|
||
|
||
if size is not None:
|
||
_reject_if_too_large(size, source="远端声明")
|
||
return size
|
||
|
||
|
||
def _content_type(headers: Any) -> Optional[str]:
|
||
raw_value = headers.get("Content-Type")
|
||
if not raw_value:
|
||
return None
|
||
media_type = raw_value.split(";", 1)[0].strip().lower()
|
||
return media_type or None
|
||
|
||
|
||
def _download(args: argparse.Namespace) -> dict[str, Any]:
|
||
output: Path = args.output
|
||
opener = urllib.request.build_opener(HTTPSOnlyRedirectHandler())
|
||
probed_size = _probe_size(opener, args.url, args.timeout)
|
||
request = urllib.request.Request(
|
||
args.url,
|
||
headers={
|
||
"Accept": "*/*",
|
||
"Accept-Encoding": "identity",
|
||
"User-Agent": USER_AGENT,
|
||
},
|
||
method="GET",
|
||
)
|
||
|
||
temp_path: Optional[Path] = None
|
||
downloaded_bytes = 0
|
||
response_size: Optional[int] = None
|
||
response_type: Optional[str] = None
|
||
try:
|
||
with tempfile.NamedTemporaryFile(
|
||
mode="wb",
|
||
prefix=f".{output.stem}.",
|
||
suffix=f".part{output.suffix}",
|
||
dir=str(output.parent),
|
||
delete=False,
|
||
) as temp_file:
|
||
temp_path = Path(temp_file.name)
|
||
with opener.open(request, timeout=args.timeout) as response:
|
||
_validate_https_url(response.geturl())
|
||
response_size = _content_length(response.headers)
|
||
response_type = _content_type(response.headers)
|
||
if response_size is not None:
|
||
_reject_if_too_large(response_size, source="下载响应声明")
|
||
|
||
while True:
|
||
chunk = response.read(CHUNK_SIZE)
|
||
if not chunk:
|
||
break
|
||
downloaded_bytes += len(chunk)
|
||
_reject_if_too_large(downloaded_bytes, source="已接收")
|
||
temp_file.write(chunk)
|
||
|
||
temp_file.flush()
|
||
os.fsync(temp_file.fileno())
|
||
|
||
if downloaded_bytes == 0:
|
||
raise ValueError("远程服务器返回了空附件")
|
||
|
||
publish_file(temp_path, output, overwrite=args.overwrite)
|
||
temp_path = None
|
||
declared_size = (
|
||
response_size if response_size is not None else probed_size
|
||
)
|
||
if response_size is not None:
|
||
size_probe = "get-content-length"
|
||
elif probed_size is not None:
|
||
size_probe = "head-content-length"
|
||
else:
|
||
size_probe = "stream"
|
||
return {
|
||
"path": str(output),
|
||
"size_bytes": downloaded_bytes,
|
||
"declared_size_bytes": declared_size,
|
||
"size_limit_bytes": MAX_ATTACHMENT_BYTES,
|
||
"size_probe": size_probe,
|
||
"content_type": response_type,
|
||
}
|
||
finally:
|
||
if temp_path is not None:
|
||
try:
|
||
temp_path.unlink(missing_ok=True)
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
def _failure_message(exc: Exception) -> str:
|
||
if isinstance(exc, urllib.error.HTTPError):
|
||
return f"附件下载失败：远程服务器返回 HTTP {exc.code}"
|
||
if isinstance(exc, (TimeoutError, socket.timeout)):
|
||
return "附件下载失败：连接或读取超时"
|
||
if isinstance(exc, urllib.error.URLError):
|
||
if isinstance(exc.reason, (TimeoutError, socket.timeout)):
|
||
return "附件下载失败：连接或读取超时"
|
||
return "附件下载失败：无法访问远程服务器"
|
||
return failure_message(exc)
|
||
|
||
|
||
def main(argv: Optional[list[str]] = None) -> int:
|
||
try:
|
||
args = _parse_args(sys.argv[1:] if argv is None else argv)
|
||
result = _download(args)
|
||
except Exception as exc:
|
||
emit({"ok": False, "error": _failure_message(exc)})
|
||
return 1
|
||
emit({"ok": True, **result})
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|