This commit is contained in:
Your Name
2026-08-26 17:18:09 +08:00
parent 327a0bc42f
commit 4ac6990efe
20 changed files with 1336 additions and 130 deletions
+82 -66
View File
@@ -31,6 +31,8 @@ import zlib
from typing import Any
from urllib.parse import urlencode
from rpa_engine.egress_channels import source_bound_requests_session
logger = logging.getLogger("douyin_im.image_upload")
_LOCAL_URL_RE = re.compile(
@@ -269,10 +271,8 @@ def _decode_sts(sts_token: str) -> tuple[str, str]:
return "", ""
def _fetch_im_upload_sts(session) -> tuple[str, str, str, str]:
def _fetch_im_upload_sts(session, source_ip: str = "") -> tuple[str, str, str, str]:
"""返回 (access_key_id, secret_access_key, sts_token, space_name)。"""
import requests
from .auth import DouyinAuth
from .dy_util import (
DEFAULT_USER_AGENT,
@@ -328,15 +328,16 @@ def _fetch_im_upload_sts(session) -> tuple[str, str, str, str]:
"Referer": "https://www.douyin.com/",
"Accept": "application/json, text/plain, */*",
}
resp = requests.get(
IM_UPLOAD_CONFIG_URL,
params=params,
headers=headers,
cookies=auth.cookie,
timeout=20,
verify=False,
proxies=_requests_proxies(),
)
with source_bound_requests_session(source_ip) as client:
resp = client.get(
IM_UPLOAD_CONFIG_URL,
params=params,
headers=headers,
cookies=auth.cookie,
timeout=20,
verify=False,
proxies=None if source_ip else _requests_proxies(),
)
data = _safe_json(resp)
if data.get("error"):
raise RuntimeError(f"获取 IM 上传配置失败:{data['error']}")
@@ -453,10 +454,8 @@ def _extract_apply_inner(data: dict[str, Any]) -> tuple[str, str, str, str]:
def _vod_apply_upload_inner(
ak: str, sk: str, token: str, space: str, file_size: int
ak: str, sk: str, token: str, space: str, file_size: int, source_ip: str = ""
) -> tuple[str, str, str, str]:
import requests
from .dy_util import DEFAULT_USER_AGENT
now = datetime.datetime.utcnow()
@@ -483,20 +482,21 @@ def _vod_apply_upload_inner(
secret_access_key=sk,
service=VOD_SERVICE,
)
resp = requests.get(
f"{VOD_HOST}?{qs}",
headers={
"accept": "*/*",
"authorization": authorization,
"user-agent": DEFAULT_USER_AGENT,
"x-amz-date": amz_date,
"x-amz-security-token": token,
"Referer": "https://www.douyin.com/",
},
timeout=30,
verify=False,
proxies=_requests_proxies(),
)
with source_bound_requests_session(source_ip) as client:
resp = client.get(
f"{VOD_HOST}?{qs}",
headers={
"accept": "*/*",
"authorization": authorization,
"user-agent": DEFAULT_USER_AGENT,
"x-amz-date": amz_date,
"x-amz-security-token": token,
"Referer": "https://www.douyin.com/",
},
timeout=30,
verify=False,
proxies=None if source_ip else _requests_proxies(),
)
data = _safe_json(resp)
if data.get("error"):
raise RuntimeError(f"申请上传地址失败:{data['error']}")
@@ -512,10 +512,14 @@ def _vod_apply_upload_inner(
# ---------------------------------------------------------------------------
def _vod_upload_binary(
host: str, store_uri: str, jwt_auth: str, user_id: str, raw: bytes, session=None
host: str,
store_uri: str,
jwt_auth: str,
user_id: str,
raw: bytes,
session=None,
source_ip: str = "",
) -> None:
import requests
from .dy_util import DEFAULT_USER_AGENT
crc32 = format(zlib.crc32(raw) & 0xFFFFFFFF, "08x")
@@ -530,14 +534,15 @@ def _vod_upload_binary(
}
if user_id:
headers["X-Storage-U"] = str(user_id)
resp = requests.post(
url,
headers=headers,
data=raw,
timeout=60,
verify=False,
proxies=_requests_proxies(),
)
with source_bound_requests_session(source_ip) as client:
resp = client.post(
url,
headers=headers,
data=raw,
timeout=60,
verify=False,
proxies=None if source_ip else _requests_proxies(),
)
data = _safe_json(resp)
if data.get("error"):
raise RuntimeError(f"上传图片数据失败:{data['error']}")
@@ -550,10 +555,8 @@ def _vod_upload_binary(
# ---------------------------------------------------------------------------
def _vod_commit_upload_inner(
ak: str, sk: str, token: str, space: str, session_key: str
ak: str, sk: str, token: str, space: str, session_key: str, source_ip: str = ""
) -> dict[str, Any]:
import requests
from .dy_util import DEFAULT_USER_AGENT
now = datetime.datetime.utcnow()
@@ -580,23 +583,24 @@ def _vod_commit_upload_inner(
signed_headers=signed_headers,
service=VOD_SERVICE,
)
resp = requests.post(
f"{VOD_HOST}?{qs}",
data=body,
headers={
"accept": "*/*",
"authorization": authorization,
"content-type": "application/json",
"user-agent": DEFAULT_USER_AGENT,
"x-amz-content-sha256": payload_hash,
"x-amz-date": amz_date,
"x-amz-security-token": token,
"Referer": "https://www.douyin.com/",
},
timeout=30,
verify=False,
proxies=_requests_proxies(),
)
with source_bound_requests_session(source_ip) as client:
resp = client.post(
f"{VOD_HOST}?{qs}",
data=body,
headers={
"accept": "*/*",
"authorization": authorization,
"content-type": "application/json",
"user-agent": DEFAULT_USER_AGENT,
"x-amz-content-sha256": payload_hash,
"x-amz-date": amz_date,
"x-amz-security-token": token,
"Referer": "https://www.douyin.com/",
},
timeout=30,
verify=False,
proxies=None if source_ip else _requests_proxies(),
)
data = _safe_json(resp)
if data.get("error"):
raise RuntimeError(f"确认上传失败:{data['error']}")
@@ -613,6 +617,7 @@ def upload_im_image(
*,
filename: str = "image.jpg",
content_type: str = "image/jpeg",
source_ip: str = "",
) -> dict[str, Any]:
"""上传图片到抖音 IM 私信图床(VOD/zhenzhen 空间)。
@@ -622,16 +627,16 @@ def upload_im_image(
if not raw:
return {"error": "图片为空"}
try:
ak, sk, token, space = _fetch_im_upload_sts(session)
ak, sk, token, space = _fetch_im_upload_sts(session, source_ip)
host, store_uri, jwt_auth, session_key = _vod_apply_upload_inner(
ak, sk, token, space, len(raw)
ak, sk, token, space, len(raw), source_ip
)
if not host or not store_uri or not jwt_auth:
return {"error": "申请上传地址失败:缺少 UploadHost/StoreUri/Auth"}
user_id = str(getattr(session, "my_uid", "") or "")
_vod_upload_binary(host, store_uri, jwt_auth, user_id, raw, session)
_vod_commit_upload_inner(ak, sk, token, space, session_key)
_vod_upload_binary(host, store_uri, jwt_auth, user_id, raw, session, source_ip)
_vod_commit_upload_inner(ak, sk, token, space, session_key, source_ip)
uri = store_uri.lstrip("/")
out: dict[str, Any] = {"uri": uri, "md5": hashlib.md5(raw).hexdigest()}
@@ -650,7 +655,12 @@ def upload_im_image(
return {"error": str(exc)}
def prepare_image_reply_spec(spec: dict[str, Any], session, upload_dir: str) -> tuple[dict[str, Any], str]:
def prepare_image_reply_spec(
spec: dict[str, Any],
session,
upload_dir: str,
source_ip: str = "",
) -> tuple[dict[str, Any], str]:
"""若图片仍是本地地址,则上传到抖音 CDN 并补全 uri。返回 (spec, error)。"""
if spec.get("type") != "image":
return spec, ""
@@ -705,7 +715,13 @@ def prepare_image_reply_spec(spec: dict[str, Any], session, upload_dir: str) ->
return spec, "图片地址必须是抖音 CDN 或本地上传后的地址,外部 URL 无法用于 IM 发送"
return spec, "缺少可上传的图片数据"
uploaded = upload_im_image(session, raw, filename=filename, content_type=content_type)
uploaded = upload_im_image(
session,
raw,
filename=filename,
content_type=content_type,
source_ip=source_ip,
)
if uploaded.get("error"):
return spec, uploaded["error"]
if not uploaded.get("uri"):