mi
473 lines
17 KiB
Python
473 lines
17 KiB
Python
import asyncio
|
|
import hashlib
|
|
import ipaddress
|
|
import json
|
|
import socket
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from typing import Any
|
|
from urllib.parse import urlparse
|
|
|
|
import boto3
|
|
import httpx
|
|
from botocore.config import Config
|
|
from redis.asyncio import Redis
|
|
|
|
from app.settings import Settings
|
|
|
|
RATE_LIMIT_LUA = """
|
|
local current = redis.call('INCR', KEYS[1])
|
|
if current == 1 then redis.call('EXPIRE', KEYS[1], ARGV[1]) end
|
|
local ttl = redis.call('TTL', KEYS[1])
|
|
return {current, ttl}
|
|
"""
|
|
|
|
|
|
class DependencyFailure(Exception):
|
|
def __init__(
|
|
self,
|
|
code: str = "dependency_unavailable",
|
|
timeout: bool = False,
|
|
*,
|
|
terminal: bool = False,
|
|
retryable: bool = True,
|
|
) -> None:
|
|
super().__init__(code)
|
|
self.code = code
|
|
self.timeout = timeout
|
|
self.terminal = terminal
|
|
self.retryable = retryable
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class CircuitBreaker:
|
|
threshold: int
|
|
open_seconds: float
|
|
failures: int = 0
|
|
opened_at: float | None = None
|
|
|
|
def allow(self) -> bool:
|
|
if self.opened_at is None:
|
|
return True
|
|
if time.monotonic() - self.opened_at >= self.open_seconds:
|
|
self.opened_at = None
|
|
self.failures = max(0, self.threshold - 1)
|
|
return True
|
|
return False
|
|
|
|
def success(self) -> None:
|
|
self.failures = 0
|
|
self.opened_at = None
|
|
|
|
def failure(self) -> None:
|
|
self.failures += 1
|
|
if self.failures >= self.threshold:
|
|
self.opened_at = time.monotonic()
|
|
|
|
|
|
class RateLimiter:
|
|
def __init__(self, redis: Redis) -> None:
|
|
self.redis = redis
|
|
|
|
async def consume(self, key: str, limit: int, window: int) -> int:
|
|
try:
|
|
count, ttl = await self.redis.eval(RATE_LIMIT_LUA, 1, key, window)
|
|
except Exception as exc:
|
|
raise DependencyFailure() from exc
|
|
if int(count) > limit:
|
|
return max(1, int(ttl))
|
|
return 0
|
|
|
|
@staticmethod
|
|
def key(identity_type: str, identity: str, route: str, window: int) -> str:
|
|
safe_identity = hashlib.sha256(identity.encode()).hexdigest()[:32]
|
|
bucket = int(time.time()) // window
|
|
return f"han:api:rl:{identity_type}:{safe_identity}:{route}:{bucket}"
|
|
|
|
|
|
class RedisIdempotency:
|
|
def __init__(self, redis: Redis) -> None:
|
|
self.redis = redis
|
|
|
|
async def get(self, scope: str, user_id: uuid.UUID, key: str) -> dict[str, Any] | None:
|
|
key_hash = hashlib.sha256(key.encode()).hexdigest()
|
|
raw = await self.redis.get(f"han:api:idem:{scope}:{user_id}:{key_hash}")
|
|
return json.loads(raw) if raw else None
|
|
|
|
async def put(self, scope: str, user_id: uuid.UUID, key: str, value: dict[str, Any]) -> None:
|
|
key_hash = hashlib.sha256(key.encode()).hexdigest()
|
|
await self.redis.set(
|
|
f"han:api:idem:{scope}:{user_id}:{key_hash}",
|
|
json.dumps(value, separators=(",", ":"), default=str),
|
|
ex=86400,
|
|
)
|
|
|
|
|
|
class SafetyClient:
|
|
def __init__(self, settings: Settings, http: httpx.AsyncClient) -> None:
|
|
self.settings = settings
|
|
self.http = http
|
|
self.breaker = CircuitBreaker(
|
|
settings.message_safety_circuit_failure_threshold,
|
|
settings.message_safety_circuit_open_sec,
|
|
)
|
|
|
|
async def check(self, payload: dict[str, Any], request_id: str) -> dict[str, Any]:
|
|
return await self._call(
|
|
"POST",
|
|
f"{self.settings.message_safety_api_prefix}/messages/check",
|
|
request_id,
|
|
json=payload,
|
|
timeout=self.settings.message_safety_post_timeout_sec,
|
|
)
|
|
|
|
async def poll(self, location: str, request_id: str) -> dict[str, Any]:
|
|
path = self._poll_path(location)
|
|
return await self._call(
|
|
"GET",
|
|
path,
|
|
request_id,
|
|
timeout=2,
|
|
)
|
|
|
|
def _poll_path(self, location: str) -> str:
|
|
expected_prefix = f"{self.settings.message_safety_api_prefix}/messages/tasks/"
|
|
parsed = urlparse(location)
|
|
if parsed.scheme or parsed.netloc or parsed.query or parsed.fragment:
|
|
raise DependencyFailure("invalid_safety_location", terminal=True, retryable=False)
|
|
if not parsed.path.startswith(expected_prefix):
|
|
raise DependencyFailure("invalid_safety_location", terminal=True, retryable=False)
|
|
task_id = parsed.path.removeprefix(expected_prefix)
|
|
try:
|
|
uuid.UUID(task_id)
|
|
except ValueError as exc:
|
|
raise DependencyFailure(
|
|
"invalid_safety_location", terminal=True, retryable=False
|
|
) from exc
|
|
return parsed.path
|
|
|
|
async def _call(self, method: str, path: str, request_id: str, **kwargs: Any) -> dict[str, Any]:
|
|
if not self.breaker.allow():
|
|
raise DependencyFailure()
|
|
headers = {
|
|
"X-Service-Token": self.settings.message_safety_service_token.get_secret_value(),
|
|
"X-Request-ID": request_id,
|
|
}
|
|
try:
|
|
response = await self.http.request(
|
|
method,
|
|
f"{str(self.settings.message_safety_url).rstrip('/')}{path}",
|
|
headers=headers,
|
|
**kwargs,
|
|
)
|
|
except httpx.TimeoutException as exc:
|
|
self.breaker.failure()
|
|
raise DependencyFailure(timeout=True) from exc
|
|
except httpx.HTTPError as exc:
|
|
self.breaker.failure()
|
|
raise DependencyFailure() from exc
|
|
if response.status_code >= 500:
|
|
self.breaker.failure()
|
|
code, terminal, retryable = "dependency_unavailable", False, True
|
|
try:
|
|
details = response.json().get("error", {}).get("details", {})
|
|
code = response.json().get("error", {}).get("code", code)
|
|
terminal = details.get("terminal") is True
|
|
retryable = details.get("retryable") is not False
|
|
except (AttributeError, ValueError):
|
|
pass
|
|
raise DependencyFailure(code, terminal=terminal, retryable=retryable)
|
|
if response.status_code not in (200, 202, 403):
|
|
if response.status_code == 401:
|
|
self.breaker.failure()
|
|
code = "safety_request_rejected"
|
|
try:
|
|
code = response.json().get("error", {}).get("code", code)
|
|
except (AttributeError, ValueError):
|
|
pass
|
|
retryable = response.status_code in (404, 429)
|
|
raise DependencyFailure(
|
|
code,
|
|
terminal=not retryable,
|
|
retryable=retryable,
|
|
)
|
|
try:
|
|
body = response.json()
|
|
except ValueError as exc:
|
|
self.breaker.failure()
|
|
raise DependencyFailure() from exc
|
|
if not isinstance(body, dict):
|
|
self.breaker.failure()
|
|
raise DependencyFailure()
|
|
status = response.status_code
|
|
expected_verdict = {200: "allow", 202: "pending", 403: "deny"}[status]
|
|
if (
|
|
body.get("verdict") != expected_verdict
|
|
or body.get("processing_mode") not in ("standard", "mock")
|
|
or type(body.get("config_version")) is not int
|
|
or not body.get("rules_version")
|
|
):
|
|
self.breaker.failure()
|
|
raise DependencyFailure("invalid_safety_response")
|
|
if status == 202:
|
|
location = response.headers.get("Location")
|
|
retry_after = response.headers.get("Retry-After")
|
|
if (
|
|
body["processing_mode"] != "standard"
|
|
or not location
|
|
or not retry_after
|
|
or not body.get("task_id")
|
|
or not body.get("expires_at")
|
|
or type(body.get("poll_after_ms")) is not int
|
|
):
|
|
self.breaker.failure()
|
|
raise DependencyFailure("invalid_safety_response")
|
|
try:
|
|
location_task_id = self._poll_path(location).rsplit("/", 1)[-1]
|
|
except DependencyFailure as exc:
|
|
self.breaker.failure()
|
|
raise DependencyFailure("invalid_safety_response") from exc
|
|
if body["task_id"] != location_task_id:
|
|
self.breaker.failure()
|
|
raise DependencyFailure("invalid_safety_response")
|
|
try:
|
|
if int(retry_after) <= 0 or body["poll_after_ms"] <= 0:
|
|
raise ValueError
|
|
datetime_value = body["expires_at"].replace("Z", "+00:00")
|
|
datetime.fromisoformat(datetime_value)
|
|
except (AttributeError, TypeError, ValueError) as exc:
|
|
self.breaker.failure()
|
|
raise DependencyFailure("invalid_safety_response") from exc
|
|
body["_location"] = location
|
|
body["_retry_after"] = retry_after
|
|
elif not body.get("rule_id") or (
|
|
status == 403 and body.get("reason_code") != "message_blocked"
|
|
):
|
|
self.breaker.failure()
|
|
raise DependencyFailure("invalid_safety_response")
|
|
self.breaker.success()
|
|
body["_status"] = response.status_code
|
|
return body
|
|
|
|
async def ready(self) -> bool:
|
|
try:
|
|
response = await self.http.get(
|
|
f"{str(self.settings.message_safety_url).rstrip('/')}/internal/safety/status",
|
|
timeout=2,
|
|
)
|
|
if response.status_code != 200:
|
|
return False
|
|
body = response.json()
|
|
capabilities = body.get("capabilities", {})
|
|
return (
|
|
body.get("status") in {"ok", "degraded"}
|
|
and body.get("processing_mode") == "standard"
|
|
and type(body.get("config_version")) is int
|
|
and all(
|
|
capabilities.get(name) == "ready"
|
|
for name in ("text", "links", "files", "worker")
|
|
)
|
|
)
|
|
except (httpx.HTTPError, AttributeError, ValueError):
|
|
return False
|
|
|
|
|
|
class OpenLinesClient:
|
|
def __init__(self, settings: Settings, http: httpx.AsyncClient) -> None:
|
|
self.settings = settings
|
|
self.http = http
|
|
self.breaker = CircuitBreaker(
|
|
settings.bitrix_local_app_circuit_failure_threshold,
|
|
settings.bitrix_local_app_circuit_open_sec,
|
|
)
|
|
|
|
async def send(
|
|
self, message_id: uuid.UUID, payload: dict[str, Any], request_id: str
|
|
) -> dict[str, Any]:
|
|
if not self.breaker.allow():
|
|
raise DependencyFailure()
|
|
try:
|
|
response = await self.http.post(
|
|
f"{str(self.settings.bitrix_local_app_base_url).rstrip('/')}"
|
|
"/internal/openlines/v1/messages",
|
|
json=payload,
|
|
headers={
|
|
"Authorization": "Bearer "
|
|
+ self.settings.bitrix_local_app_internal_token.get_secret_value(),
|
|
"Idempotency-Key": str(message_id),
|
|
"X-Request-ID": request_id,
|
|
},
|
|
timeout=self.settings.bitrix_local_app_http_timeout_sec,
|
|
)
|
|
response.raise_for_status()
|
|
except httpx.TimeoutException as exc:
|
|
self.breaker.failure()
|
|
raise DependencyFailure(timeout=True) from exc
|
|
except httpx.HTTPError as exc:
|
|
self.breaker.failure()
|
|
raise DependencyFailure() from exc
|
|
self.breaker.success()
|
|
return response.json()
|
|
|
|
async def ready(self) -> bool:
|
|
try:
|
|
response = await self.http.get(
|
|
f"{str(self.settings.bitrix_local_app_base_url).rstrip('/')}"
|
|
"/internal/openlines/v1/status",
|
|
headers={
|
|
"Authorization": "Bearer "
|
|
+ self.settings.bitrix_local_app_internal_token.get_secret_value()
|
|
},
|
|
timeout=2,
|
|
)
|
|
return response.is_success
|
|
except httpx.HTTPError:
|
|
return False
|
|
|
|
|
|
async def fresh_openlines_payload(
|
|
payload: dict[str, Any], s3: "S3Client"
|
|
) -> dict[str, Any]:
|
|
result = json.loads(json.dumps(payload, default=str))
|
|
for file in result.get("message", {}).get("files", []):
|
|
bucket = file.pop("_storage_bucket")
|
|
key = file.pop("_object_key")
|
|
file["download_url"] = await s3.presign_get(bucket, key)
|
|
return result
|
|
|
|
|
|
class S3Client:
|
|
def __init__(self, settings: Settings) -> None:
|
|
self.settings = settings
|
|
self.client = boto3.client(
|
|
"s3",
|
|
endpoint_url=str(settings.selectel_s3_endpoint_url),
|
|
aws_access_key_id=settings.selectel_s3_access_key.get_secret_value(),
|
|
aws_secret_access_key=settings.selectel_s3_secret_key.get_secret_value(),
|
|
config=Config(
|
|
signature_version="s3v4",
|
|
connect_timeout=3,
|
|
read_timeout=10,
|
|
retries={"max_attempts": 2},
|
|
s3={"addressing_style": "virtual"},
|
|
),
|
|
)
|
|
|
|
async def ready(self) -> bool:
|
|
try:
|
|
for bucket in (
|
|
self.settings.selectel_s3_bucket_quarantine,
|
|
self.settings.selectel_s3_bucket_attachments,
|
|
self.settings.selectel_s3_bucket_documents,
|
|
):
|
|
await asyncio.to_thread(self.client.head_bucket, Bucket=bucket)
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
async def presign_put(self, key: str, mime: str, ttl: int) -> str:
|
|
return await asyncio.to_thread(
|
|
self.client.generate_presigned_url,
|
|
"put_object",
|
|
Params={
|
|
"Bucket": self.settings.selectel_s3_bucket_quarantine,
|
|
"Key": key,
|
|
"ContentType": mime,
|
|
},
|
|
ExpiresIn=ttl,
|
|
)
|
|
|
|
async def presign_get(self, bucket: str, key: str, ttl: int = 300) -> str:
|
|
return await asyncio.to_thread(
|
|
self.client.generate_presigned_url,
|
|
"get_object",
|
|
Params={"Bucket": bucket, "Key": key},
|
|
ExpiresIn=ttl,
|
|
)
|
|
|
|
async def head(self, bucket: str, key: str) -> dict[str, Any]:
|
|
return await asyncio.to_thread(self.client.head_object, Bucket=bucket, Key=key)
|
|
|
|
async def promote(
|
|
self,
|
|
source_key: str,
|
|
destination_key: str,
|
|
*,
|
|
version_id: str,
|
|
etag: str,
|
|
) -> None:
|
|
await asyncio.to_thread(
|
|
self.client.copy_object,
|
|
Bucket=self.settings.selectel_s3_bucket_attachments,
|
|
Key=destination_key,
|
|
CopySource={
|
|
"Bucket": self.settings.selectel_s3_bucket_quarantine,
|
|
"Key": source_key,
|
|
"VersionId": version_id,
|
|
},
|
|
CopySourceIfMatch=etag,
|
|
)
|
|
# Keep the immutable source version until quarantine lifecycle expiry.
|
|
# A crash after copy but before the DB checkpoint can then safely retry
|
|
# the same conditional copy without losing its source.
|
|
|
|
async def delete_quarantine(self, key: str) -> None:
|
|
await asyncio.to_thread(
|
|
self.client.delete_object,
|
|
Bucket=self.settings.selectel_s3_bucket_quarantine,
|
|
Key=key,
|
|
)
|
|
|
|
async def delete(self, bucket: str, key: str) -> None:
|
|
await asyncio.to_thread(self.client.delete_object, Bucket=bucket, Key=key)
|
|
|
|
async def upload_inbound(
|
|
self,
|
|
http: httpx.AsyncClient,
|
|
download_url: str,
|
|
destination_key: str,
|
|
mime_type: str,
|
|
max_bytes: int,
|
|
) -> tuple[int, str]:
|
|
parsed = urlparse(download_url)
|
|
if parsed.scheme != "https" or not parsed.hostname:
|
|
raise DependencyFailure("unsafe_inbound_url")
|
|
try:
|
|
addresses = await asyncio.to_thread(
|
|
socket.getaddrinfo, parsed.hostname, parsed.port or 443, type=socket.SOCK_STREAM
|
|
)
|
|
except OSError as exc:
|
|
raise DependencyFailure("unsafe_inbound_url") from exc
|
|
if any(
|
|
ipaddress.ip_address(address[4][0]).is_private
|
|
or ipaddress.ip_address(address[4][0]).is_loopback
|
|
or ipaddress.ip_address(address[4][0]).is_link_local
|
|
or ipaddress.ip_address(address[4][0]).is_reserved
|
|
for address in addresses
|
|
):
|
|
raise DependencyFailure("unsafe_inbound_url")
|
|
data = bytearray()
|
|
try:
|
|
async with http.stream(
|
|
"GET", download_url, timeout=10, follow_redirects=False
|
|
) as response:
|
|
response.raise_for_status()
|
|
if response.headers.get("content-type", "").split(";")[0] != mime_type:
|
|
raise DependencyFailure("inbound_mime_mismatch")
|
|
async for chunk in response.aiter_bytes():
|
|
data.extend(chunk)
|
|
if len(data) > max_bytes:
|
|
raise DependencyFailure("inbound_file_too_large")
|
|
except httpx.HTTPError as exc:
|
|
raise DependencyFailure() from exc
|
|
await asyncio.to_thread(
|
|
self.client.put_object,
|
|
Bucket=self.settings.selectel_s3_bucket_attachments,
|
|
Key=destination_key,
|
|
Body=bytes(data),
|
|
ContentType=mime_type,
|
|
)
|
|
return len(data), hashlib.sha256(data).hexdigest()
|