72 lines
1.8 KiB
Python
72 lines
1.8 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from yuxi.utils.logging_config import logger
|
||
|
|
|
||
|
|
_UNSAFE_TRANSPORT_FIELDS = frozenset(
|
||
|
|
{
|
||
|
|
"agent",
|
||
|
|
"cert",
|
||
|
|
"cert_file",
|
||
|
|
"key",
|
||
|
|
"key_file",
|
||
|
|
"dispatcher",
|
||
|
|
"proxy",
|
||
|
|
"session",
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
_AUTH_RESPONSE_MAX_BYTES = 1 * 1024 * 1024
|
||
|
|
|
||
|
|
|
||
|
|
def sanitize_google_auth_init(kwargs: dict[str, Any]) -> dict[str, Any]:
|
||
|
|
cleaned: dict[str, Any] = {}
|
||
|
|
removed: list[str] = []
|
||
|
|
for key, value in kwargs.items():
|
||
|
|
if key in _UNSAFE_TRANSPORT_FIELDS:
|
||
|
|
removed.append(key)
|
||
|
|
continue
|
||
|
|
cleaned[key] = value
|
||
|
|
if removed:
|
||
|
|
logger.warning(f"SSRF guard: removed unsafe transport fields: {removed}")
|
||
|
|
return cleaned
|
||
|
|
|
||
|
|
|
||
|
|
def sanitize_request_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
|
||
|
|
cleaned: dict[str, Any] = {}
|
||
|
|
removed: list[str] = []
|
||
|
|
for key, value in kwargs.items():
|
||
|
|
if key in _UNSAFE_TRANSPORT_FIELDS:
|
||
|
|
removed.append(key)
|
||
|
|
continue
|
||
|
|
cleaned[key] = value
|
||
|
|
if removed:
|
||
|
|
logger.warning(f"SSRF guard: removed unsafe request fields: {removed}")
|
||
|
|
return cleaned
|
||
|
|
|
||
|
|
|
||
|
|
def apply_ssrf_guard() -> dict[str, Any]:
|
||
|
|
return {
|
||
|
|
"agent": None,
|
||
|
|
"cert": None,
|
||
|
|
"cert_file": None,
|
||
|
|
"key": None,
|
||
|
|
"key_file": None,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def check_auth_response_size(data: bytes) -> bytes:
|
||
|
|
if len(data) > _AUTH_RESPONSE_MAX_BYTES:
|
||
|
|
logger.warning(
|
||
|
|
f"Google Auth response exceeds size limit: {len(data)} bytes > {_AUTH_RESPONSE_MAX_BYTES} bytes, truncating"
|
||
|
|
)
|
||
|
|
return data[:_AUTH_RESPONSE_MAX_BYTES]
|
||
|
|
return data
|
||
|
|
|
||
|
|
|
||
|
|
def check_auth_response_text(text: str) -> str:
|
||
|
|
data = text.encode("utf-8", errors="replace")
|
||
|
|
truncated = check_auth_response_size(data)
|
||
|
|
return truncated.decode("utf-8", errors="replace")
|