mirror of
https://github.com/fromchat-messenger/web.git
synced 2026-09-22 19:15:08 +03:00
Add streaming encryption support
This commit is contained in:
@@ -537,6 +537,54 @@ async def complete_resumable_upload(upload_id: str, request: Request):
|
|||||||
return await complete_resumable_upload_internal(upload_id, int(user_id_header))
|
return await complete_resumable_upload_internal(upload_id, int(user_id_header))
|
||||||
|
|
||||||
|
|
||||||
|
async def get_resumable_upload_blob_path_internal(upload_id: str, user_id: int) -> dict:
|
||||||
|
"""Return on-disk path to completed resumable ciphertext (no base64)."""
|
||||||
|
meta = _read_resumable_meta(upload_id)
|
||||||
|
_assert_resumable_access(meta, user_id)
|
||||||
|
if not meta.get("complete"):
|
||||||
|
raise HTTPException(status_code=409, detail="Upload not completed")
|
||||||
|
data_path = _resumable_data_path(upload_id)
|
||||||
|
if not data_path.exists():
|
||||||
|
raise HTTPException(status_code=404, detail="Upload payload not found")
|
||||||
|
return {
|
||||||
|
"upload_id": upload_id,
|
||||||
|
"filename": meta["filename"],
|
||||||
|
"file_size": int(meta.get("total_size", 0)),
|
||||||
|
"encrypted_file_path": str(data_path.resolve()),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
async def upload_encrypted_file_from_path_internal(
|
||||||
|
filename: str,
|
||||||
|
source_path: Path,
|
||||||
|
content_type: str = "application/octet-stream",
|
||||||
|
allowed_user_ids: list[int] | None = None,
|
||||||
|
) -> dict:
|
||||||
|
"""Store a pre-encrypted file by copying from a local path (no base64)."""
|
||||||
|
allowed_user_ids = allowed_user_ids or []
|
||||||
|
src = Path(source_path)
|
||||||
|
if not src.is_file():
|
||||||
|
raise HTTPException(status_code=400, detail="source_path is not a file")
|
||||||
|
_ensure_dirs()
|
||||||
|
original_name = _secure_filename(filename or "file")
|
||||||
|
uid = uuid.uuid4().hex
|
||||||
|
stored_name = f"{uid}_{original_name}"
|
||||||
|
dest = FILES_DIR / stored_name
|
||||||
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
import shutil
|
||||||
|
|
||||||
|
shutil.copyfile(src, dest)
|
||||||
|
dest.chmod(0o600)
|
||||||
|
_store_file_permissions(stored_name, allowed_user_ids)
|
||||||
|
size = dest.stat().st_size
|
||||||
|
return {
|
||||||
|
"file_id": stored_name,
|
||||||
|
"filename": original_name,
|
||||||
|
"size": size,
|
||||||
|
"path": f"/uploads/files/encrypted/{stored_name}",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
async def get_resumable_upload_data_internal(upload_id: str, user_id: int) -> dict:
|
async def get_resumable_upload_data_internal(upload_id: str, user_id: int) -> dict:
|
||||||
"""Internal implementation for in-process calls."""
|
"""Internal implementation for in-process calls."""
|
||||||
meta = _read_resumable_meta(upload_id)
|
meta = _read_resumable_meta(upload_id)
|
||||||
|
|||||||
@@ -34,7 +34,8 @@ from ..service_calls import (
|
|||||||
get_resumable_upload_status_in_storage,
|
get_resumable_upload_status_in_storage,
|
||||||
upload_resumable_chunk_in_storage,
|
upload_resumable_chunk_in_storage,
|
||||||
complete_resumable_upload_in_storage,
|
complete_resumable_upload_in_storage,
|
||||||
get_resumable_upload_data_in_storage,
|
get_resumable_upload_blob_path_in_storage,
|
||||||
|
store_encrypted_file_from_path,
|
||||||
delete_resumable_upload_in_storage,
|
delete_resumable_upload_in_storage,
|
||||||
)
|
)
|
||||||
from .messaging import messagingManager, convert_dm_envelope, convert_dm_envelope_for_user
|
from .messaging import messagingManager, convert_dm_envelope, convert_dm_envelope_for_user
|
||||||
@@ -262,10 +263,12 @@ async def send_encrypted_message(
|
|||||||
]
|
]
|
||||||
|
|
||||||
for upload_id in request.uploaded_file_ids:
|
for upload_id in request.uploaded_file_ids:
|
||||||
uploaded_payload = await get_resumable_upload_data_in_storage(upload_id, current_user.id)
|
uploaded_payload = await get_resumable_upload_blob_path_in_storage(
|
||||||
|
upload_id, current_user.id
|
||||||
|
)
|
||||||
all_transport_files.append(
|
all_transport_files.append(
|
||||||
{
|
{
|
||||||
"encrypted_file_data_b64": uploaded_payload["encrypted_file_data_b64"],
|
"encrypted_file_path": uploaded_payload["encrypted_file_path"],
|
||||||
"filename": uploaded_payload["filename"],
|
"filename": uploaded_payload["filename"],
|
||||||
"file_size": uploaded_payload["file_size"],
|
"file_size": uploaded_payload["file_size"],
|
||||||
"upload_id": upload_id,
|
"upload_id": upload_id,
|
||||||
@@ -280,10 +283,17 @@ async def send_encrypted_message(
|
|||||||
sender_public_key_b64=request.sender_public_key_b64,
|
sender_public_key_b64=request.sender_public_key_b64,
|
||||||
recipient_public_key_b64=request.recipient_public_key_b64,
|
recipient_public_key_b64=request.recipient_public_key_b64,
|
||||||
transport_files=[
|
transport_files=[
|
||||||
{
|
(
|
||||||
"encrypted_file_data_b64": str(f["encrypted_file_data_b64"]),
|
{
|
||||||
"filename": str(f.get("filename", "file")),
|
"encrypted_file_path": str(f["encrypted_file_path"]),
|
||||||
}
|
"filename": str(f.get("filename", "file")),
|
||||||
|
}
|
||||||
|
if f.get("encrypted_file_path")
|
||||||
|
else {
|
||||||
|
"encrypted_file_data_b64": str(f["encrypted_file_data_b64"]),
|
||||||
|
"filename": str(f.get("filename", "file")),
|
||||||
|
}
|
||||||
|
)
|
||||||
for f in all_transport_files
|
for f in all_transport_files
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
@@ -319,13 +329,27 @@ async def send_encrypted_message(
|
|||||||
|
|
||||||
for i, tf in enumerate(all_transport_files):
|
for i, tf in enumerate(all_transport_files):
|
||||||
fr = file_results[i]
|
fr = file_results[i]
|
||||||
file_storage_result = await store_encrypted_file(
|
ciphertext_path = fr.get("ciphertext_path")
|
||||||
encrypted_file_data_b64=fr["ciphertext"],
|
if ciphertext_path:
|
||||||
filename=str(tf["filename"]),
|
file_storage_result = await store_encrypted_file_from_path(
|
||||||
content_type="application/octet-stream",
|
source_path=str(ciphertext_path),
|
||||||
sender_id=current_user.id,
|
filename=str(tf["filename"]),
|
||||||
recipient_id=request.recipient_id,
|
content_type="application/octet-stream",
|
||||||
)
|
sender_id=current_user.id,
|
||||||
|
recipient_id=request.recipient_id,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
Path(ciphertext_path).unlink(missing_ok=True)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
file_storage_result = await store_encrypted_file(
|
||||||
|
encrypted_file_data_b64=fr["ciphertext"],
|
||||||
|
filename=str(tf["filename"]),
|
||||||
|
content_type="application/octet-stream",
|
||||||
|
sender_id=current_user.id,
|
||||||
|
recipient_id=request.recipient_id,
|
||||||
|
)
|
||||||
|
|
||||||
df = DMFile(
|
df = DMFile(
|
||||||
message_id=dm_envelope.id,
|
message_id=dm_envelope.id,
|
||||||
|
|||||||
@@ -564,6 +564,57 @@ async def complete_resumable_upload_in_storage(
|
|||||||
return r.json()
|
return r.json()
|
||||||
|
|
||||||
|
|
||||||
|
async def get_resumable_upload_blob_path_in_storage(
|
||||||
|
upload_id: str,
|
||||||
|
user_id: int,
|
||||||
|
timeout: float = 30.0,
|
||||||
|
) -> Dict[str, Any]:
|
||||||
|
mod = _get_file_storage_module()
|
||||||
|
if mod:
|
||||||
|
try:
|
||||||
|
return await mod.get_resumable_upload_blob_path_internal(upload_id, user_id)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("In-process file_storage.get_resumable_upload_blob_path failed: %s", e)
|
||||||
|
raise
|
||||||
|
|
||||||
|
file_storage_url = os.getenv("FILE_STORAGE_URL") or os.getenv("FILE_STORAGE_SERVICE_URL") or _default_file_storage_base_url()
|
||||||
|
url = f"{file_storage_url.rstrip('/')}/uploads/resumable/{upload_id}/blob-path"
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||||
|
r = await client.get(url, headers={"X-User-ID": str(user_id)})
|
||||||
|
r.raise_for_status()
|
||||||
|
return r.json()
|
||||||
|
|
||||||
|
|
||||||
|
async def store_encrypted_file_from_path(
|
||||||
|
source_path: str,
|
||||||
|
filename: str,
|
||||||
|
content_type: str = "application/octet-stream",
|
||||||
|
sender_id: int = None,
|
||||||
|
recipient_id: int = None,
|
||||||
|
timeout: float = 120.0,
|
||||||
|
) -> Dict[str, Any]:
|
||||||
|
mod = _get_file_storage_module()
|
||||||
|
allowed_user_ids: list[int] = []
|
||||||
|
if sender_id is not None:
|
||||||
|
allowed_user_ids.append(sender_id)
|
||||||
|
if recipient_id is not None:
|
||||||
|
allowed_user_ids.append(recipient_id)
|
||||||
|
|
||||||
|
if mod:
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
return await mod.upload_encrypted_file_from_path_internal(
|
||||||
|
filename=filename,
|
||||||
|
source_path=Path(source_path),
|
||||||
|
content_type=content_type,
|
||||||
|
allowed_user_ids=allowed_user_ids,
|
||||||
|
)
|
||||||
|
|
||||||
|
raise RuntimeError("store_encrypted_file_from_path requires in-process file_storage")
|
||||||
|
|
||||||
|
|
||||||
async def get_resumable_upload_data_in_storage(
|
async def get_resumable_upload_data_in_storage(
|
||||||
upload_id: str,
|
upload_id: str,
|
||||||
user_id: int,
|
user_id: int,
|
||||||
|
|||||||
@@ -11,6 +11,8 @@ Handles:
|
|||||||
import os
|
import os
|
||||||
import base64
|
import base64
|
||||||
import logging
|
import logging
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import BinaryIO
|
||||||
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
||||||
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
|
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
|
||||||
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
|
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
|
||||||
@@ -25,6 +27,13 @@ TRANSPORT_NONCE_SIZE = 24 # For X25519 transport encryption (PyNaCl Box/XSalsa2
|
|||||||
MEK_NONCE_SIZE = 12 # For AES-GCM content encryption
|
MEK_NONCE_SIZE = 12 # For AES-GCM content encryption
|
||||||
MEK_SIZE = 32 # Message Encryption Key size
|
MEK_SIZE = 32 # Message Encryption Key size
|
||||||
|
|
||||||
|
# Client streaming transport format (chunked AES-256-GCM): FCAE | version | frames…
|
||||||
|
FCAE_MAGIC = b"FCAE"
|
||||||
|
FCAE_VERSION = 1
|
||||||
|
FCAE_PREFIX_BYTES = len(FCAE_MAGIC) + 1
|
||||||
|
FCAE_FRAME_LENGTH_BYTES = 4
|
||||||
|
TRANSPORT_FILE_KEY_CONTEXT = "fromchat_transport_file_v1"
|
||||||
|
|
||||||
|
|
||||||
def generate_mek() -> bytes:
|
def generate_mek() -> bytes:
|
||||||
"""Generate a random Message Encryption Key (32 bytes)."""
|
"""Generate a random Message Encryption Key (32 bytes)."""
|
||||||
@@ -123,6 +132,106 @@ def decrypt_transport_message(
|
|||||||
raise
|
raise
|
||||||
|
|
||||||
|
|
||||||
|
def is_fcae_transport_blob(prefix: bytes) -> bool:
|
||||||
|
return len(prefix) >= len(FCAE_MAGIC) and prefix[: len(FCAE_MAGIC)] == FCAE_MAGIC
|
||||||
|
|
||||||
|
|
||||||
|
def derive_transport_file_aes_key(
|
||||||
|
client_public_key_b64: str,
|
||||||
|
ephemeral_private_key: X25519PrivateKey,
|
||||||
|
) -> bytes:
|
||||||
|
client_public_bytes = base64.b64decode(client_public_key_b64)
|
||||||
|
server_private_bytes = ephemeral_private_key.private_bytes_raw()
|
||||||
|
shared = sodium.crypto_box_beforenm(client_public_bytes, server_private_bytes)
|
||||||
|
return derive_key_from_shared_secret(shared, TRANSPORT_FILE_KEY_CONTEXT)
|
||||||
|
|
||||||
|
|
||||||
|
def _read_fcae_frame_payload(source: BinaryIO) -> tuple[bytes, bytes] | None:
|
||||||
|
length_bytes = source.read(FCAE_FRAME_LENGTH_BYTES)
|
||||||
|
if not length_bytes:
|
||||||
|
return None
|
||||||
|
if len(length_bytes) < FCAE_FRAME_LENGTH_BYTES:
|
||||||
|
raise ValueError("Truncated FCAE frame length")
|
||||||
|
frame_len = int.from_bytes(length_bytes, byteorder="big", signed=False)
|
||||||
|
if frame_len <= MEK_NONCE_SIZE:
|
||||||
|
raise ValueError("Invalid FCAE frame length")
|
||||||
|
frame = source.read(frame_len)
|
||||||
|
if len(frame) < frame_len:
|
||||||
|
raise ValueError("Truncated FCAE frame")
|
||||||
|
iv = frame[:MEK_NONCE_SIZE]
|
||||||
|
ciphertext = frame[MEK_NONCE_SIZE:]
|
||||||
|
return iv, ciphertext
|
||||||
|
|
||||||
|
|
||||||
|
def _decrypt_fcae_transport_stream_io(
|
||||||
|
client_public_key_b64: str,
|
||||||
|
source: BinaryIO,
|
||||||
|
ephemeral_private_key: X25519PrivateKey,
|
||||||
|
) -> bytes:
|
||||||
|
prefix = source.read(FCAE_PREFIX_BYTES)
|
||||||
|
if len(prefix) < FCAE_PREFIX_BYTES:
|
||||||
|
raise ValueError("FCAE blob is too short")
|
||||||
|
if not is_fcae_transport_blob(prefix):
|
||||||
|
raise ValueError("Not an FCAE transport blob")
|
||||||
|
if prefix[4] != FCAE_VERSION:
|
||||||
|
raise ValueError("Unsupported FCAE version")
|
||||||
|
aes_key = derive_transport_file_aes_key(client_public_key_b64, ephemeral_private_key)
|
||||||
|
cipher = AESGCM(aes_key)
|
||||||
|
parts: list[bytes] = []
|
||||||
|
while True:
|
||||||
|
frame = _read_fcae_frame_payload(source)
|
||||||
|
if frame is None:
|
||||||
|
break
|
||||||
|
iv, ciphertext = frame
|
||||||
|
parts.append(cipher.decrypt(iv, ciphertext, None))
|
||||||
|
return b"".join(parts)
|
||||||
|
|
||||||
|
|
||||||
|
def decrypt_fcae_transport_blob_to_file(
|
||||||
|
client_public_key_b64: str,
|
||||||
|
encrypted_path: Path,
|
||||||
|
ephemeral_private_key: X25519PrivateKey,
|
||||||
|
output_path: Path,
|
||||||
|
) -> int:
|
||||||
|
"""Stream-decrypt FCAE transport ciphertext from disk to a plaintext file."""
|
||||||
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
total_out = 0
|
||||||
|
with open(encrypted_path, "rb") as enc, open(output_path, "wb") as out:
|
||||||
|
prefix = enc.read(FCAE_PREFIX_BYTES)
|
||||||
|
if not is_fcae_transport_blob(prefix):
|
||||||
|
raise ValueError("Not an FCAE transport blob")
|
||||||
|
if prefix[4] != FCAE_VERSION:
|
||||||
|
raise ValueError("Unsupported FCAE version")
|
||||||
|
aes_key = derive_transport_file_aes_key(client_public_key_b64, ephemeral_private_key)
|
||||||
|
cipher = AESGCM(aes_key)
|
||||||
|
while True:
|
||||||
|
frame = _read_fcae_frame_payload(enc)
|
||||||
|
if frame is None:
|
||||||
|
break
|
||||||
|
iv, ciphertext = frame
|
||||||
|
plain = cipher.decrypt(iv, ciphertext, None)
|
||||||
|
out.write(plain)
|
||||||
|
total_out += len(plain)
|
||||||
|
return total_out
|
||||||
|
|
||||||
|
|
||||||
|
def encrypt_message_to_file(plaintext_path: Path, mek: bytes, output_path: Path) -> str:
|
||||||
|
"""AES-GCM encrypt a file on disk; returns nonce_b64. Ciphertext written to output_path."""
|
||||||
|
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
||||||
|
|
||||||
|
nonce = generate_nonce(MEK_NONCE_SIZE)
|
||||||
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
encryptor = Cipher(algorithms.AES(mek), modes.GCM(nonce)).encryptor()
|
||||||
|
with open(plaintext_path, "rb") as src, open(output_path, "wb") as dst:
|
||||||
|
while True:
|
||||||
|
chunk = src.read(1024 * 1024)
|
||||||
|
if not chunk:
|
||||||
|
break
|
||||||
|
dst.write(encryptor.update(chunk))
|
||||||
|
dst.write(encryptor.finalize())
|
||||||
|
return base64.b64encode(nonce).decode("utf-8")
|
||||||
|
|
||||||
|
|
||||||
def decrypt_transport_blob(
|
def decrypt_transport_blob(
|
||||||
client_public_key_b64: str,
|
client_public_key_b64: str,
|
||||||
encrypted_blob: bytes,
|
encrypted_blob: bytes,
|
||||||
@@ -146,6 +255,15 @@ def decrypt_transport_blob(
|
|||||||
Returns:
|
Returns:
|
||||||
Decrypted plaintext bytes.
|
Decrypted plaintext bytes.
|
||||||
"""
|
"""
|
||||||
|
if is_fcae_transport_blob(encrypted_blob):
|
||||||
|
import io
|
||||||
|
|
||||||
|
return _decrypt_fcae_transport_stream_io(
|
||||||
|
client_public_key_b64,
|
||||||
|
io.BytesIO(encrypted_blob),
|
||||||
|
ephemeral_private_key,
|
||||||
|
)
|
||||||
|
|
||||||
if len(encrypted_blob) < nonce_size + 16:
|
if len(encrypted_blob) < nonce_size + 16:
|
||||||
# crypto_box has a MAC; ciphertext must have at least some overhead.
|
# crypto_box has a MAC; ciphertext must have at least some overhead.
|
||||||
raise ValueError("Encrypted blob is too short to contain nonce + ciphertext")
|
raise ValueError("Encrypted blob is too short to contain nonce + ciphertext")
|
||||||
|
|||||||
@@ -15,7 +15,9 @@ import sys
|
|||||||
import time
|
import time
|
||||||
import base64
|
import base64
|
||||||
import os
|
import os
|
||||||
from typing import Dict, Any
|
import tempfile
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Dict, Any, Union
|
||||||
from fastapi import FastAPI, HTTPException, status
|
from fastapi import FastAPI, HTTPException, status
|
||||||
from nacl.exceptions import CryptoError
|
from nacl.exceptions import CryptoError
|
||||||
from fastapi.middleware.cors import CORSMiddleware
|
from fastapi.middleware.cors import CORSMiddleware
|
||||||
@@ -27,7 +29,14 @@ logger = logging.getLogger("uvicorn.error")
|
|||||||
_B64_DECODE_KW = {"validate": True} if sys.version_info >= (3, 11) else {}
|
_B64_DECODE_KW = {"validate": True} if sys.version_info >= (3, 11) else {}
|
||||||
|
|
||||||
# Import encryption modules
|
# Import encryption modules
|
||||||
from .encryption import generate_nonce, TRANSPORT_NONCE_SIZE, decrypt_transport_blob, decrypt_transport_message
|
from .encryption import (
|
||||||
|
generate_nonce,
|
||||||
|
TRANSPORT_NONCE_SIZE,
|
||||||
|
decrypt_transport_blob,
|
||||||
|
decrypt_transport_message,
|
||||||
|
is_fcae_transport_blob,
|
||||||
|
decrypt_fcae_transport_blob_to_file,
|
||||||
|
)
|
||||||
from .processor import process_encrypted_message, process_encrypted_message_and_files
|
from .processor import process_encrypted_message, process_encrypted_message_and_files
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -374,40 +383,84 @@ async def process_message_with_files(
|
|||||||
)
|
)
|
||||||
|
|
||||||
plaintext_files: list[bytes] = []
|
plaintext_files: list[bytes] = []
|
||||||
|
plaintext_file_paths: list[Path | None] = []
|
||||||
filenames: list[str] = []
|
filenames: list[str] = []
|
||||||
for idx, tf in enumerate(transport_files):
|
temp_paths: list[Path] = []
|
||||||
enc_b64 = tf.get("encrypted_file_data_b64", "")
|
try:
|
||||||
try:
|
for idx, tf in enumerate(transport_files):
|
||||||
transport_blob = base64.b64decode(enc_b64, **_B64_DECODE_KW)
|
enc_path = (tf.get("encrypted_file_path") or "").strip()
|
||||||
except Exception as e:
|
if enc_path:
|
||||||
logger.error("Invalid base64 for transport file index=%s filename=%r: %s", idx, tf.get("filename"), e)
|
blob_path = Path(enc_path)
|
||||||
raise
|
if not blob_path.is_file():
|
||||||
try:
|
raise ValueError(f"Transport file path missing index={idx}")
|
||||||
plaintext_files.append(
|
prefix = blob_path.read_bytes()[: len(b"FCAE") + 1]
|
||||||
decrypt_transport_blob(
|
if is_fcae_transport_blob(prefix):
|
||||||
client_public_key_b64=client_public_key_b64,
|
plain_tmp = Path(tempfile.mkstemp(prefix="fcae-plain-", suffix=".bin")[1])
|
||||||
encrypted_blob=transport_blob,
|
temp_paths.append(plain_tmp)
|
||||||
ephemeral_private_key=private_key,
|
decrypt_fcae_transport_blob_to_file(
|
||||||
)
|
client_public_key_b64=client_public_key_b64,
|
||||||
)
|
encrypted_path=blob_path,
|
||||||
except Exception as e:
|
ephemeral_private_key=private_key,
|
||||||
logger.error(
|
output_path=plain_tmp,
|
||||||
"Transport file decrypt failed index=%s filename=%r (check same ephemeral as message): %s",
|
)
|
||||||
idx,
|
plaintext_files.append(b"")
|
||||||
tf.get("filename"),
|
plaintext_file_paths.append(plain_tmp)
|
||||||
e,
|
else:
|
||||||
)
|
transport_blob = blob_path.read_bytes()
|
||||||
raise
|
plaintext_files.append(
|
||||||
filenames.append(tf.get("filename", "file"))
|
decrypt_transport_blob(
|
||||||
|
client_public_key_b64=client_public_key_b64,
|
||||||
|
encrypted_blob=transport_blob,
|
||||||
|
ephemeral_private_key=private_key,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
plaintext_file_paths.append(None)
|
||||||
|
else:
|
||||||
|
enc_b64 = tf.get("encrypted_file_data_b64", "")
|
||||||
|
try:
|
||||||
|
transport_blob = base64.b64decode(enc_b64, **_B64_DECODE_KW)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(
|
||||||
|
"Invalid base64 for transport file index=%s filename=%r: %s",
|
||||||
|
idx,
|
||||||
|
tf.get("filename"),
|
||||||
|
e,
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
try:
|
||||||
|
plaintext_files.append(
|
||||||
|
decrypt_transport_blob(
|
||||||
|
client_public_key_b64=client_public_key_b64,
|
||||||
|
encrypted_blob=transport_blob,
|
||||||
|
ephemeral_private_key=private_key,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(
|
||||||
|
"Transport file decrypt failed index=%s filename=%r (check same ephemeral as message): %s",
|
||||||
|
idx,
|
||||||
|
tf.get("filename"),
|
||||||
|
e,
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
plaintext_file_paths.append(None)
|
||||||
|
filenames.append(tf.get("filename", "file"))
|
||||||
|
|
||||||
return process_encrypted_message_and_files(
|
return process_encrypted_message_and_files(
|
||||||
plaintext_message=plaintext_message,
|
plaintext_message=plaintext_message,
|
||||||
plaintext_files=plaintext_files,
|
plaintext_files=plaintext_files,
|
||||||
filenames=filenames,
|
filenames=filenames,
|
||||||
compliance_public_key_b64=compliance_public_key_b64,
|
plaintext_file_paths=plaintext_file_paths,
|
||||||
sender_public_key_b64=sender_public_key_b64,
|
compliance_public_key_b64=compliance_public_key_b64,
|
||||||
recipient_public_key_b64=recipient_public_key_b64,
|
sender_public_key_b64=sender_public_key_b64,
|
||||||
)
|
recipient_public_key_b64=recipient_public_key_b64,
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
for p in temp_paths:
|
||||||
|
try:
|
||||||
|
p.unlink(missing_ok=True)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
@app.post("/process-with-files", response_model=None)
|
@app.post("/process-with-files", response_model=None)
|
||||||
|
|||||||
@@ -22,6 +22,7 @@ from .encryption import (
|
|||||||
decrypt_transport_message,
|
decrypt_transport_message,
|
||||||
generate_mek,
|
generate_mek,
|
||||||
encrypt_message,
|
encrypt_message,
|
||||||
|
encrypt_message_to_file,
|
||||||
wrap_mek,
|
wrap_mek,
|
||||||
derive_shared_secret,
|
derive_shared_secret,
|
||||||
derive_key_from_shared_secret,
|
derive_key_from_shared_secret,
|
||||||
@@ -205,6 +206,9 @@ def _generate_thumbnail(image_bytes: bytes) -> tuple[str | None, list[int]]:
|
|||||||
return (None, [1, 1])
|
return (None, [1, 1])
|
||||||
|
|
||||||
|
|
||||||
|
_LARGE_FILE_THUMB_BYTES = 32 * 1024 * 1024
|
||||||
|
|
||||||
|
|
||||||
def process_encrypted_message_and_files(
|
def process_encrypted_message_and_files(
|
||||||
plaintext_message: bytes,
|
plaintext_message: bytes,
|
||||||
plaintext_files: list[bytes],
|
plaintext_files: list[bytes],
|
||||||
@@ -212,6 +216,7 @@ def process_encrypted_message_and_files(
|
|||||||
compliance_public_key_b64: str,
|
compliance_public_key_b64: str,
|
||||||
sender_public_key_b64: str,
|
sender_public_key_b64: str,
|
||||||
recipient_public_key_b64: str,
|
recipient_public_key_b64: str,
|
||||||
|
plaintext_file_paths: list[Path | None] | None = None,
|
||||||
) -> Dict[str, Any]:
|
) -> Dict[str, Any]:
|
||||||
"""
|
"""
|
||||||
Process a message and its attached files using a single MEK.
|
Process a message and its attached files using a single MEK.
|
||||||
@@ -230,6 +235,10 @@ def process_encrypted_message_and_files(
|
|||||||
if len(filenames) != len(plaintext_files):
|
if len(filenames) != len(plaintext_files):
|
||||||
filenames = [f"file_{i}" for i in range(len(plaintext_files))]
|
filenames = [f"file_{i}" for i in range(len(plaintext_files))]
|
||||||
|
|
||||||
|
paths = plaintext_file_paths or [None] * len(plaintext_files)
|
||||||
|
if len(paths) < len(plaintext_files):
|
||||||
|
paths = paths + [None] * (len(plaintext_files) - len(paths))
|
||||||
|
|
||||||
# One MEK for everything in this envelope
|
# One MEK for everything in this envelope
|
||||||
mek = generate_mek()
|
mek = generate_mek()
|
||||||
|
|
||||||
@@ -239,11 +248,24 @@ def process_encrypted_message_and_files(
|
|||||||
file_sizes: list[int] = []
|
file_sizes: list[int] = []
|
||||||
for i, f_bytes in enumerate(plaintext_files):
|
for i, f_bytes in enumerate(plaintext_files):
|
||||||
name = filenames[i] if i < len(filenames) else ""
|
name = filenames[i] if i < len(filenames) else ""
|
||||||
file_sizes.append(len(f_bytes))
|
path = paths[i]
|
||||||
if Path(name).suffix.lower() in _IMAGE_EXTENSIONS:
|
size = int(path.stat().st_size) if path is not None else len(f_bytes)
|
||||||
|
file_sizes.append(size)
|
||||||
|
if (
|
||||||
|
path is None
|
||||||
|
and Path(name).suffix.lower() in _IMAGE_EXTENSIONS
|
||||||
|
):
|
||||||
thumb_b64, wh = _generate_thumbnail(f_bytes)
|
thumb_b64, wh = _generate_thumbnail(f_bytes)
|
||||||
file_thumbnails.append(thumb_b64 or "")
|
file_thumbnails.append(thumb_b64 or "")
|
||||||
file_aspect_ratios.append(wh)
|
file_aspect_ratios.append(wh)
|
||||||
|
elif (
|
||||||
|
path is not None
|
||||||
|
and size <= _LARGE_FILE_THUMB_BYTES
|
||||||
|
and Path(name).suffix.lower() in _IMAGE_EXTENSIONS
|
||||||
|
):
|
||||||
|
thumb_b64, wh = _generate_thumbnail(path.read_bytes())
|
||||||
|
file_thumbnails.append(thumb_b64 or "")
|
||||||
|
file_aspect_ratios.append(wh)
|
||||||
else:
|
else:
|
||||||
file_thumbnails.append("")
|
file_thumbnails.append("")
|
||||||
file_aspect_ratios.append([1, 1])
|
file_aspect_ratios.append([1, 1])
|
||||||
@@ -270,10 +292,17 @@ def process_encrypted_message_and_files(
|
|||||||
|
|
||||||
# Encrypt files (same MEK, per-file nonce)
|
# Encrypt files (same MEK, per-file nonce)
|
||||||
files_out: list[Dict[str, Any]] = []
|
files_out: list[Dict[str, Any]] = []
|
||||||
|
import tempfile
|
||||||
|
|
||||||
for i, f_bytes in enumerate(plaintext_files):
|
for i, f_bytes in enumerate(plaintext_files):
|
||||||
f_nonce, f_ciphertext = encrypt_message(f_bytes, mek)
|
path = paths[i]
|
||||||
entry: Dict[str, Any] = {"nonce": f_nonce, "ciphertext": f_ciphertext}
|
if path is not None:
|
||||||
files_out.append(entry)
|
enc_tmp = Path(tempfile.mkstemp(prefix="mek-enc-", suffix=".bin")[1])
|
||||||
|
f_nonce = encrypt_message_to_file(path, mek, enc_tmp)
|
||||||
|
files_out.append({"nonce": f_nonce, "ciphertext_path": str(enc_tmp)})
|
||||||
|
else:
|
||||||
|
f_nonce, f_ciphertext = encrypt_message(f_bytes, mek)
|
||||||
|
files_out.append({"nonce": f_nonce, "ciphertext": f_ciphertext})
|
||||||
|
|
||||||
# Derive wrap keys deterministically (same as existing flow)
|
# Derive wrap keys deterministically (same as existing flow)
|
||||||
sender_key_bytes = base64.b64decode(sender_public_key_b64)
|
sender_key_bytes = base64.b64decode(sender_public_key_b64)
|
||||||
|
|||||||
Reference in New Issue
Block a user