317 lines
9.3 KiB
Python
317 lines
9.3 KiB
Python
import contextvars
|
|
import hashlib
|
|
import json
|
|
import threading
|
|
import time
|
|
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
|
|
|
|
|
SENSITIVE_HEADERS = {
|
|
"authorization",
|
|
"cookie",
|
|
"proxy-authorization",
|
|
"set-cookie",
|
|
"x-api-key",
|
|
"api-key",
|
|
}
|
|
SENSITIVE_QUERY_KEYS = {
|
|
"access_key",
|
|
"access_token",
|
|
"api_key",
|
|
"apikey",
|
|
"authorization",
|
|
"key",
|
|
"secret",
|
|
"signature",
|
|
"token",
|
|
}
|
|
|
|
_current_prompt_id = contextvars.ContextVar("account_manager_prompt_id", default=None)
|
|
_fallback_prompt_id = None
|
|
_fallback_lock = threading.Lock()
|
|
_recorder = None
|
|
_installed = False
|
|
|
|
MEDIA_CONTENT_TYPES = (
|
|
"image/",
|
|
"video/",
|
|
"audio/",
|
|
"application/octet-stream",
|
|
)
|
|
MEDIA_KEY_PARTS = (
|
|
"image",
|
|
"video",
|
|
"audio",
|
|
"base64",
|
|
"binary",
|
|
"file_data",
|
|
"first_frame",
|
|
"last_frame",
|
|
"mask",
|
|
)
|
|
LARGE_MEDIA_STRING_BYTES = 64 * 1024
|
|
|
|
|
|
def set_current_job(prompt_id: str) -> None:
|
|
global _fallback_prompt_id
|
|
prompt_id = str(prompt_id or "") or None
|
|
_current_prompt_id.set(prompt_id)
|
|
with _fallback_lock:
|
|
_fallback_prompt_id = prompt_id
|
|
|
|
|
|
def clear_current_job(prompt_id: str = None) -> None:
|
|
global _fallback_prompt_id
|
|
current = _current_prompt_id.get()
|
|
if prompt_id is None or current == prompt_id:
|
|
_current_prompt_id.set(None)
|
|
with _fallback_lock:
|
|
if prompt_id is None or _fallback_prompt_id == prompt_id:
|
|
_fallback_prompt_id = None
|
|
|
|
|
|
def current_job() -> str | None:
|
|
prompt_id = _current_prompt_id.get()
|
|
if prompt_id:
|
|
return prompt_id
|
|
with _fallback_lock:
|
|
return _fallback_prompt_id
|
|
|
|
|
|
def _redact_url(url) -> str:
|
|
text = str(url or "")
|
|
try:
|
|
parts = urlsplit(text)
|
|
query = urlencode(
|
|
[
|
|
(key, "[REDACTED]" if key.lower() in SENSITIVE_QUERY_KEYS else value)
|
|
for key, value in parse_qsl(parts.query, keep_blank_values=True)
|
|
],
|
|
doseq=True,
|
|
)
|
|
return urlunsplit((parts.scheme, parts.netloc, parts.path, query, parts.fragment))
|
|
except ValueError:
|
|
return text
|
|
|
|
|
|
def _headers(value) -> str:
|
|
if not value:
|
|
return "{}"
|
|
items = value.items() if hasattr(value, "items") else value
|
|
return json.dumps(
|
|
{
|
|
str(key): "[REDACTED]" if str(key).lower() in SENSITIVE_HEADERS else str(item)
|
|
for key, item in items
|
|
},
|
|
ensure_ascii=False,
|
|
separators=(",", ":"),
|
|
)
|
|
|
|
|
|
def _body(value) -> bytes | None:
|
|
if value is None or value is False:
|
|
return None
|
|
if isinstance(value, bytes):
|
|
return value
|
|
if isinstance(value, bytearray):
|
|
return bytes(value)
|
|
if isinstance(value, memoryview):
|
|
return value.tobytes()
|
|
if isinstance(value, str):
|
|
return value.encode("utf-8")
|
|
return json.dumps(value, ensure_ascii=False, separators=(",", ":"), default=str).encode(
|
|
"utf-8"
|
|
)
|
|
|
|
|
|
def _content_type(headers) -> str:
|
|
if not headers:
|
|
return ""
|
|
try:
|
|
value = headers.get("content-type") or headers.get("Content-Type")
|
|
except AttributeError:
|
|
value = None
|
|
for key, item in headers:
|
|
if str(key).lower() == "content-type":
|
|
value = item
|
|
break
|
|
return str(value or "").split(";", 1)[0].strip().lower()
|
|
|
|
|
|
def _body_reference(raw: bytes, content_type: str, reason: str) -> bytes:
|
|
return json.dumps(
|
|
{
|
|
"$account_manager_body_ref": {
|
|
"stored": False,
|
|
"reason": reason,
|
|
"content_type": content_type or None,
|
|
"content_bytes": len(raw),
|
|
"sha256": hashlib.sha256(raw).hexdigest(),
|
|
}
|
|
},
|
|
ensure_ascii=False,
|
|
separators=(",", ":"),
|
|
).encode("utf-8")
|
|
|
|
|
|
def _media_string_reference(value: str, reason: str) -> dict:
|
|
raw = value.encode("utf-8")
|
|
return {
|
|
"$account_manager_media_ref": {
|
|
"stored": False,
|
|
"reason": reason,
|
|
"encoded_bytes": len(raw),
|
|
"sha256": hashlib.sha256(raw).hexdigest(),
|
|
}
|
|
}
|
|
|
|
|
|
def _strip_embedded_media(value, key: str = ""):
|
|
if isinstance(value, dict):
|
|
return {
|
|
item_key: _strip_embedded_media(item, str(item_key).lower())
|
|
for item_key, item in value.items()
|
|
}
|
|
if isinstance(value, list):
|
|
return [_strip_embedded_media(item, key) for item in value]
|
|
if not isinstance(value, str):
|
|
return value
|
|
|
|
lowered = value[:64].lower()
|
|
if lowered.startswith(("data:image/", "data:video/", "data:audio/")):
|
|
return _media_string_reference(value, "embedded_media_data_uri")
|
|
if len(value.encode("utf-8")) >= LARGE_MEDIA_STRING_BYTES and any(
|
|
part in key for part in MEDIA_KEY_PARTS
|
|
):
|
|
return _media_string_reference(value, "embedded_media_field")
|
|
return value
|
|
|
|
|
|
def _audit_body(value, headers=None) -> bytes | None:
|
|
raw = _body(value)
|
|
if raw is None:
|
|
return None
|
|
content_type = _content_type(headers)
|
|
if content_type.startswith(MEDIA_CONTENT_TYPES):
|
|
return _body_reference(raw, content_type, "binary_media_body")
|
|
if content_type == "multipart/form-data":
|
|
return _body_reference(raw, content_type, "multipart_body")
|
|
if content_type == "application/json" or (
|
|
not content_type and raw.lstrip().startswith((b"{", b"["))
|
|
):
|
|
try:
|
|
parsed = json.loads(raw)
|
|
except (UnicodeDecodeError, json.JSONDecodeError):
|
|
return raw
|
|
stripped = _strip_embedded_media(parsed)
|
|
return json.dumps(
|
|
stripped, ensure_ascii=False, separators=(",", ":"), default=str
|
|
).encode("utf-8")
|
|
return raw
|
|
|
|
|
|
def _record(client, request, response=None, error: BaseException = None, *, stream=False) -> None:
|
|
prompt_id = current_job()
|
|
if not prompt_id or _recorder is None:
|
|
return
|
|
try:
|
|
response_body = None
|
|
# Reading .content consumes the raw stream used by COS and other downloaders.
|
|
if response is not None and not stream:
|
|
try:
|
|
response_body = _audit_body(response.content, response.headers)
|
|
except Exception:
|
|
response_body = _audit_body(
|
|
getattr(response, "_content", None),
|
|
getattr(response, "headers", None),
|
|
)
|
|
_recorder(
|
|
prompt_id=prompt_id,
|
|
recorded_at=time.time(),
|
|
client=client,
|
|
method=str(getattr(request, "method", "") or ""),
|
|
url=_redact_url(getattr(request, "url", "")),
|
|
request_headers=_headers(getattr(request, "headers", None)),
|
|
request_body=_audit_body(
|
|
getattr(request, "body", None) or getattr(request, "content", None),
|
|
getattr(request, "headers", None),
|
|
),
|
|
response_status=(getattr(response, "status_code", None) if response is not None else None),
|
|
response_headers=_headers(getattr(response, "headers", None)),
|
|
response_body=response_body,
|
|
error=str(error) if error is not None else None,
|
|
)
|
|
except Exception:
|
|
return
|
|
|
|
|
|
def _fallback_request(method, url, kwargs):
|
|
class Request:
|
|
pass
|
|
|
|
request = Request()
|
|
request.method = method
|
|
request.url = url
|
|
request.headers = kwargs.get("headers")
|
|
request.body = kwargs.get("json", kwargs.get("data"))
|
|
return request
|
|
|
|
|
|
def install_api_audit(recorder) -> None:
|
|
global _installed, _recorder
|
|
_recorder = recorder
|
|
if _installed:
|
|
return
|
|
|
|
try:
|
|
import requests
|
|
|
|
original = requests.sessions.Session.request
|
|
|
|
def requests_request(session, method, url, **kwargs):
|
|
try:
|
|
response = original(session, method, url, **kwargs)
|
|
except Exception as error:
|
|
_record("requests", _fallback_request(method, url, kwargs), error=error)
|
|
raise
|
|
stream = kwargs.get("stream")
|
|
if stream is None:
|
|
stream = session.stream
|
|
_record("requests", response.request, response, stream=stream)
|
|
return response
|
|
|
|
requests.sessions.Session.request = requests_request
|
|
except ImportError:
|
|
pass
|
|
|
|
try:
|
|
import httpx
|
|
|
|
original_sync = httpx.Client.request
|
|
original_async = httpx.AsyncClient.request
|
|
|
|
def httpx_request(client, method, url, *args, **kwargs):
|
|
try:
|
|
response = original_sync(client, method, url, *args, **kwargs)
|
|
except Exception as error:
|
|
_record("httpx", _fallback_request(method, url, kwargs), error=error)
|
|
raise
|
|
_record("httpx", response.request, response)
|
|
return response
|
|
|
|
async def httpx_async_request(client, method, url, *args, **kwargs):
|
|
try:
|
|
response = await original_async(client, method, url, *args, **kwargs)
|
|
except Exception as error:
|
|
_record("httpx-async", _fallback_request(method, url, kwargs), error=error)
|
|
raise
|
|
_record("httpx-async", response.request, response)
|
|
return response
|
|
|
|
httpx.Client.request = httpx_request
|
|
httpx.AsyncClient.request = httpx_async_request
|
|
except ImportError:
|
|
pass
|
|
|
|
_installed = True
|