101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
import asyncio
|
|
import hashlib
|
|
from dataclasses import dataclass
|
|
|
|
import google_crc32c
|
|
from google.api_core.exceptions import NotFound, PreconditionFailed
|
|
from google.cloud import storage
|
|
|
|
from app.config import get_settings
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StoredMime:
|
|
bucket: str
|
|
name: str
|
|
generation: int
|
|
size: int
|
|
crc32c: str
|
|
|
|
|
|
class RawMimeStorage:
|
|
def __init__(self, bucket: str, prefix: str = "raw", project: str | None = None):
|
|
self.bucket_name = bucket
|
|
self.prefix = prefix.strip("/")
|
|
self.client = storage.Client(project=project)
|
|
|
|
async def put(self, raw: bytes, delivery_key: str) -> StoredMime:
|
|
return await asyncio.to_thread(self._put, raw, delivery_key)
|
|
|
|
def _put(self, raw: bytes, delivery_key: str) -> StoredMime:
|
|
name = f"{self.prefix}/{delivery_key[:2]}/{delivery_key}.eml"
|
|
blob = self.client.bucket(self.bucket_name).blob(name)
|
|
checksum = google_crc32c.Checksum(raw).hexdigest().decode("ascii")
|
|
try:
|
|
blob.upload_from_string(
|
|
raw,
|
|
content_type="message/rfc822",
|
|
checksum="crc32c",
|
|
if_generation_match=0,
|
|
)
|
|
blob.reload()
|
|
except PreconditionFailed:
|
|
blob.reload()
|
|
return StoredMime(
|
|
bucket=self.bucket_name,
|
|
name=name,
|
|
generation=int(blob.generation),
|
|
size=len(raw),
|
|
crc32c=checksum,
|
|
)
|
|
|
|
async def get(self, reference: dict) -> bytes:
|
|
return await asyncio.to_thread(self._get, reference)
|
|
|
|
def _get(self, reference: dict) -> bytes:
|
|
generation = int(reference["generation"])
|
|
blob = self.client.bucket(reference["bucket"]).blob(
|
|
reference["name"],
|
|
generation=generation,
|
|
)
|
|
return blob.download_as_bytes(if_generation_match=generation)
|
|
|
|
async def delete(self, reference: dict) -> None:
|
|
await asyncio.to_thread(self._delete, reference)
|
|
|
|
def _delete(self, reference: dict) -> None:
|
|
generation = int(reference["generation"])
|
|
blob = self.client.bucket(reference["bucket"]).blob(
|
|
reference["name"],
|
|
generation=generation,
|
|
)
|
|
try:
|
|
blob.delete(if_generation_match=generation)
|
|
except NotFound:
|
|
# A prior delete may have removed GCS before Mongo acknowledged it.
|
|
pass
|
|
|
|
|
|
def delivery_digest(raw: bytes, recipients: list[str]) -> str:
|
|
digest = hashlib.sha256()
|
|
digest.update(raw)
|
|
digest.update(b"\0")
|
|
digest.update("\0".join(sorted(set(recipients))).encode())
|
|
return digest.hexdigest()
|
|
|
|
|
|
_storage: RawMimeStorage | None = None
|
|
|
|
|
|
def get_raw_mime_storage() -> RawMimeStorage:
|
|
global _storage
|
|
if _storage is None:
|
|
settings = get_settings()
|
|
if not settings.gcs_raw_mime_bucket:
|
|
raise RuntimeError("GCS_RAW_MIME_BUCKET is required for inbound mail.")
|
|
_storage = RawMimeStorage(
|
|
settings.gcs_raw_mime_bucket,
|
|
settings.gcs_raw_mime_prefix,
|
|
settings.gcs_project,
|
|
)
|
|
return _storage
|