# SPDX-License-Identifier: AGPL-3.0-only
# Copyright (c) 2026 Xessenger contributors
"""Private R2 media storage using the same authenticated BM02 envelope as the app.

Only this storage boundary is custom; Synapse and the network bridges are unmodified.
No persistent local media copy is required. Plaintext temp files close after use.
"""
import hashlib
import json
import os
from pathlib import Path
import re
import tempfile
import time

import boto3
from botocore.config import Config
from botocore.exceptions import ClientError
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes
from synapse.media.media_storage import FileResponder
from synapse.media.storage_provider import StorageProvider
from synapse.module_api import NOT_SPAM
from twisted.python.threadpool import ThreadPool

MAX_BYTES = 20 * 1024 * 1024
CHUNK = 256 * 1024

class EncryptedR2(StorageProvider):
    def __init__(self, hs, config):
        self.hs = hs
        self.api = hs.get_module_api()
        self.config = config
        self.master = bytes.fromhex(config["encryption_key"])
        self.bucket = config["bucket"]
        self.prefix = config["prefix"]
        self.cache_directory = hs.config.media.media_store_path
        self.paths = {}
        self.pool = ThreadPool(minthreads=1, maxthreads=4, name="xessenger-r2")
        self.pool.start()
        hs.get_reactor().addSystemEventTrigger("before", "shutdown", self.pool.stop)
        self.client = boto3.client("s3", endpoint_url=config["endpoint"], region_name="auto",
            aws_access_key_id=config["access_key_id"], aws_secret_access_key=config["secret_access_key"],
            config=Config(signature_version="s3v4", s3={"addressing_style":"path"},
                request_checksum_calculation="when_required", response_checksum_validation="when_required",
                connect_timeout=10, read_timeout=30, retries={"max_attempts":2}))
        # In Synapse 1.162, local-media-disabled uploads use a temporary file,
        # but store_file receives only the logical path. The supported media
        # check hook exposes the upload path; observing it does not bypass any
        # other checks. This also works when a later release sets upload_path.
        self.api.register_spam_checker_callbacks(check_media_file_for_spam=self.capture_path)

    async def capture_path(self, file, file_info):
        now = time.monotonic()
        self.paths = {key:value for key,value in self.paths.items() if now-value[0] < 120}
        if len(self.paths) >= 128:
            raise RuntimeError("Media staging capacity exceeded")
        self.paths[id(file_info)] = (now, file.path)
        return NOT_SPAM

    def object_key(self, path):
        if not re.fullmatch(r"[A-Za-z0-9_./:-]+", path) or path.startswith("/") or ".." in path.split("/"):
            raise ValueError("Invalid media path")
        return self.prefix + path

    def file_key(self, path):
        return HKDF(algorithm=hashes.SHA256(), length=32, salt=self.prefix.encode(),
            info=("matrix-media-v1:"+path).encode()).derive(self.master)

    async def store_file(self, path, file_info):
        captured = self.paths.pop(id(file_info), None)
        source = getattr(file_info, "upload_path", None) or (captured[1] if captured else os.path.join(self.cache_directory,path))
        await self.api.defer_to_threadpool(self.pool, self.store, path, source)

    def store(self, path, source):
        key = self.object_key(path)
        iv = os.urandom(12)
        encryptor = Cipher(algorithms.AES(self.file_key(path)), modes.GCM(iv)).encryptor()
        encryptor.authenticate_additional_data(key.encode())
        with open(source,"rb") as plain, tempfile.TemporaryFile() as encrypted:
            encrypted.write(b"BM02"+iv+b"\0"*16)
            size = 0
            while block := plain.read(CHUNK):
                size += len(block)
                if size > MAX_BYTES: raise ValueError("Media exceeds configured size limit")
                encrypted.write(encryptor.update(block))
            encrypted.write(encryptor.finalize())
            encrypted.seek(16); encrypted.write(encryptor.tag); encrypted.seek(0)
            self.client.upload_fileobj(encrypted, self.bucket, key, ExtraArgs={"ContentType":"application/octet-stream"})

    async def fetch(self, path, file_info):
        file = await self.api.defer_to_threadpool(self.pool, self.download, path)
        return FileResponder(self.hs,file) if file is not None else None

    def download(self, path):
        key = self.object_key(path)
        try: response = self.client.get_object(Bucket=self.bucket,Key=key)
        except ClientError as error:
            if error.response["Error"]["Code"] in ["NoSuchKey","404"]: return None
            raise
        stream = response["Body"]
        output = tempfile.TemporaryFile()
        try:
            if not 32 <= response.get("ContentLength",0) <= MAX_BYTES+32:
                raise ValueError("Invalid encrypted media size")
            header = stream.read(32)
            if len(header)!=32 or header[:4]!=b"BM02": raise ValueError("Invalid encrypted media")
            decryptor = Cipher(algorithms.AES(self.file_key(path)),modes.GCM(header[4:16],header[16:32])).decryptor()
            decryptor.authenticate_additional_data(key.encode())
            size = 0
            while block := stream.read(CHUNK):
                size += len(block)
                if size > MAX_BYTES: raise ValueError("Invalid encrypted media size")
                output.write(decryptor.update(block))
            output.write(decryptor.finalize())
            output.seek(0)
            return output
        except Exception:
            output.close()
            raise
        finally: stream.close()

    @staticmethod
    def parse_config(config):
        result = json.loads(Path(config["credentials_file"]).read_text())
        if not re.fullmatch(r"[a-f0-9]{64}",result.get("encryption_key","")):
            raise ValueError("Missing media encryption key")
        if not re.fullmatch(r"[A-Za-z0-9_/-]+/",result.get("prefix","")) or ".." in result["prefix"]:
            raise ValueError("Invalid media namespace")
        if not result.get("endpoint","").startswith("https://"):
            raise ValueError("Media endpoint must use HTTPS")
        return result