115 lines
3.0 KiB
Python
115 lines
3.0 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
import logging
|
||
|
|
import os
|
||
|
|
import tempfile
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
MEDIA_UPLOAD_URL = "https://oapi.dingtalk.com/media/upload"
|
||
|
|
MESSAGE_FILES_DOWNLOAD_URL = "https://api.dingtalk.com/v1.0/robot/messageFiles/download"
|
||
|
|
|
||
|
|
|
||
|
|
async def upload_media(
|
||
|
|
http: httpx.AsyncClient,
|
||
|
|
token: str,
|
||
|
|
file_path: str,
|
||
|
|
media_type: str,
|
||
|
|
) -> str | None:
|
||
|
|
if file_path.startswith("file://"):
|
||
|
|
file_path = file_path[7:]
|
||
|
|
|
||
|
|
if file_path.startswith(("http://", "https://")):
|
||
|
|
file_path = await _download_to_tmp(http, file_path)
|
||
|
|
|
||
|
|
if not os.path.exists(file_path):
|
||
|
|
logger.error("Media file not found: %s", file_path)
|
||
|
|
return None
|
||
|
|
|
||
|
|
with open(file_path, "rb") as f:
|
||
|
|
resp = await http.post(
|
||
|
|
MEDIA_UPLOAD_URL,
|
||
|
|
params={"access_token": token, "type": media_type},
|
||
|
|
files={"media": (os.path.basename(file_path), f)},
|
||
|
|
timeout=60.0,
|
||
|
|
)
|
||
|
|
data = resp.json()
|
||
|
|
if data.get("errcode") == 0:
|
||
|
|
return data.get("media_id")
|
||
|
|
logger.error("Media upload failed: %s", data)
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
async def download_image(
|
||
|
|
http: httpx.AsyncClient,
|
||
|
|
token: str,
|
||
|
|
download_code: str,
|
||
|
|
robot_code: str,
|
||
|
|
save_dir: str | None = None,
|
||
|
|
) -> str | None:
|
||
|
|
resp = await http.post(
|
||
|
|
MESSAGE_FILES_DOWNLOAD_URL,
|
||
|
|
headers={
|
||
|
|
"x-acs-dingtalk-access-token": token,
|
||
|
|
"Content-Type": "application/json",
|
||
|
|
},
|
||
|
|
json={
|
||
|
|
"downloadCode": download_code,
|
||
|
|
"robotCode": robot_code,
|
||
|
|
},
|
||
|
|
timeout=30.0,
|
||
|
|
)
|
||
|
|
data = resp.json()
|
||
|
|
download_url = data.get("downloadUrl")
|
||
|
|
if not download_url:
|
||
|
|
logger.error("Failed to get download URL: %s", data)
|
||
|
|
return None
|
||
|
|
|
||
|
|
img_resp = await http.get(download_url, timeout=30.0)
|
||
|
|
img_resp.raise_for_status()
|
||
|
|
|
||
|
|
img_data = img_resp.content
|
||
|
|
ext = _guess_extension(img_data)
|
||
|
|
filename = f"{hashlib.md5(img_data).hexdigest()}{ext}"
|
||
|
|
|
||
|
|
if save_dir:
|
||
|
|
os.makedirs(save_dir, exist_ok=True)
|
||
|
|
save_path = os.path.join(save_dir, filename)
|
||
|
|
else:
|
||
|
|
tmp = tempfile.NamedTemporaryFile(delete=False, suffix=ext)
|
||
|
|
save_path = tmp.name
|
||
|
|
tmp.close()
|
||
|
|
|
||
|
|
with open(save_path, "wb") as f:
|
||
|
|
f.write(img_data)
|
||
|
|
|
||
|
|
logger.debug("Image downloaded: %s -> %s", download_code, save_path)
|
||
|
|
return save_path
|
||
|
|
|
||
|
|
|
||
|
|
async def _download_to_tmp(http: httpx.AsyncClient, url: str) -> str:
|
||
|
|
resp = await http.get(url, timeout=60.0)
|
||
|
|
resp.raise_for_status()
|
||
|
|
data = resp.content
|
||
|
|
ext = _guess_extension(data)
|
||
|
|
|
||
|
|
tmp = tempfile.NamedTemporaryFile(delete=False, suffix=ext)
|
||
|
|
tmp.write(data)
|
||
|
|
tmp.close()
|
||
|
|
return tmp.name
|
||
|
|
|
||
|
|
|
||
|
|
def _guess_extension(data: bytes) -> str:
|
||
|
|
if data[:4] == b"\x89PNG":
|
||
|
|
return ".png"
|
||
|
|
if data[:2] == b"\xff\xd8":
|
||
|
|
return ".jpg"
|
||
|
|
if data[:4] == b"GIF8":
|
||
|
|
return ".gif"
|
||
|
|
if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
|
||
|
|
return ".webp"
|
||
|
|
return ".bin"
|