mi
344 lines
12 KiB
Python
344 lines
12 KiB
Python
import asyncio
|
|
import hashlib
|
|
import ipaddress
|
|
import json
|
|
import socket
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
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) -> None:
|
|
self.code = code
|
|
self.timeout = timeout
|
|
|
|
|
|
@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",
|
|
"/internal/safety/v1/messages/check",
|
|
request_id,
|
|
json=payload,
|
|
timeout=self.settings.message_safety_post_timeout_sec,
|
|
)
|
|
|
|
async def poll(self, task_id: str, request_id: str) -> dict[str, Any]:
|
|
return await self._call(
|
|
"GET",
|
|
f"/internal/safety/v1/messages/tasks/{task_id}",
|
|
request_id,
|
|
timeout=2,
|
|
)
|
|
|
|
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 == 401 or response.status_code >= 500:
|
|
self.breaker.failure()
|
|
raise DependencyFailure()
|
|
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()
|
|
self.breaker.success()
|
|
body["_status"] = response.status_code
|
|
return body
|
|
|
|
|
|
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(
|
|
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) -> 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,
|
|
},
|
|
)
|
|
await asyncio.to_thread(
|
|
self.client.delete_object,
|
|
Bucket=self.settings.selectel_s3_bucket_quarantine,
|
|
Key=source_key,
|
|
)
|
|
|
|
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 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()
|