feat: encrypted file transfer with propose/accept flow
New commands: /send <filepath>, /accept, /reject Protocol: - Sender proposes file (name, size, SHA-256 hash) - Recipient sees offer and chooses /accept or /reject - On accept: file chunked (64KB), encrypted with room key, sent over WebSocket - On receive: chunks reassembled, SHA-256 verified, saved to ./downloads/ - Server never sees file content (E2E encrypted, same as messages) Limits: 50MB max file size. Files saved with collision-safe naming. No server changes — server remains a dumb encrypted relay. All 79 existing tests pass. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
65ee9dee16
commit
70ddca8a1f
|
|
@ -1,8 +1,11 @@
|
|||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import ssl
|
||||
import base64
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from uuid import uuid4
|
||||
|
||||
import srp
|
||||
import requests
|
||||
|
|
@ -15,6 +18,17 @@ from rich.panel import Panel
|
|||
|
||||
srp.rfc5054_enable()
|
||||
|
||||
MAX_FILE_SIZE = 50 * 1024 * 1024 # 50 MB
|
||||
CHUNK_SIZE = 64 * 1024 # 64 KB
|
||||
|
||||
|
||||
def _human_size(size: int) -> str:
|
||||
for unit in ("B", "KB", "MB", "GB"):
|
||||
if size < 1024:
|
||||
return f"{size:.1f} {unit}"
|
||||
size /= 1024
|
||||
return f"{size:.1f} TB"
|
||||
|
||||
|
||||
class Client:
|
||||
def __init__(
|
||||
|
|
@ -42,6 +56,14 @@ class Client:
|
|||
self.connected = False
|
||||
self.running = False
|
||||
|
||||
# File transfer state
|
||||
self.pending_offer: Optional[dict] = None
|
||||
self.active_send: Optional[dict] = None # {id, path} for outgoing
|
||||
self.received_chunks: dict[str, list[bytes]] = {}
|
||||
self.transfer_meta: dict[str, dict] = {} # id -> offer metadata
|
||||
self.download_dir = Path("./downloads")
|
||||
self._ws_ref: Optional[object] = None # WebSocket reference for sending from input_loop
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
scheme = "http" if self.no_tls else "https"
|
||||
|
|
@ -175,7 +197,181 @@ class Client:
|
|||
self.console.print("[dim italic]No messages yet...[/]")
|
||||
|
||||
self.console.print("-" * 60)
|
||||
self.console.print("[dim]Type message and press Enter. 'q' to quit.[/]")
|
||||
|
||||
if self.pending_offer:
|
||||
name = self.pending_offer.get("name", "?")
|
||||
size = _human_size(self.pending_offer.get("size", 0))
|
||||
sender = self._safe_username(self.pending_offer.get("from", "?"))
|
||||
self.console.print(
|
||||
f"[yellow bold]{sender}[/] wants to send [bold]{name}[/] ({size})"
|
||||
f" — [green]/accept[/] or [red]/reject[/]"
|
||||
)
|
||||
self.console.print("-" * 60)
|
||||
|
||||
self.console.print(
|
||||
"[dim]/send <file> | /accept | /reject | q to quit[/]"
|
||||
)
|
||||
|
||||
# ── File Transfer: Sending ───────────────────────────────────────────
|
||||
|
||||
async def _send_offer(self, ws, filepath: str) -> None:
|
||||
path = Path(filepath).expanduser().resolve()
|
||||
if not path.is_file():
|
||||
self.error(f"File not found: {filepath}")
|
||||
return
|
||||
|
||||
size = path.stat().st_size
|
||||
if size > MAX_FILE_SIZE:
|
||||
self.error(f"File too large ({_human_size(size)}). Max is {_human_size(MAX_FILE_SIZE)}")
|
||||
return
|
||||
|
||||
sha = hashlib.sha256()
|
||||
with open(path, "rb") as f:
|
||||
while chunk := f.read(CHUNK_SIZE):
|
||||
sha.update(chunk)
|
||||
|
||||
transfer_id = str(uuid4())
|
||||
self.active_send = {"id": transfer_id, "path": str(path)}
|
||||
|
||||
offer = json.dumps({
|
||||
"_ft": "offer",
|
||||
"id": transfer_id,
|
||||
"name": path.name,
|
||||
"size": size,
|
||||
"sha256": sha.hexdigest(),
|
||||
})
|
||||
encrypted = self.room_fernet.encrypt(offer.encode()).decode()
|
||||
await ws.send(encrypted)
|
||||
self.info(f"Offered {path.name} ({_human_size(size)}) — waiting for accept...")
|
||||
|
||||
async def _send_file_chunks(self, ws, transfer_id: str, filepath: str) -> None:
|
||||
path = Path(filepath)
|
||||
total = path.stat().st_size
|
||||
sent = 0
|
||||
seq = 0
|
||||
|
||||
with open(path, "rb") as f:
|
||||
while chunk := f.read(CHUNK_SIZE):
|
||||
msg = json.dumps({
|
||||
"_ft": "chunk",
|
||||
"id": transfer_id,
|
||||
"seq": seq,
|
||||
"data": base64.b64encode(chunk).decode(),
|
||||
})
|
||||
encrypted = self.room_fernet.encrypt(msg.encode()).decode()
|
||||
await ws.send(encrypted)
|
||||
sent += len(chunk)
|
||||
seq += 1
|
||||
pct = int(sent * 100 / total) if total else 100
|
||||
self.console.print(f"\r[cyan]Sending: {pct}% ({_human_size(sent)}/{_human_size(total)})[/]", end="")
|
||||
await asyncio.sleep(0.01) # yield to event loop
|
||||
|
||||
done_msg = json.dumps({"_ft": "done", "id": transfer_id})
|
||||
encrypted = self.room_fernet.encrypt(done_msg.encode()).decode()
|
||||
await ws.send(encrypted)
|
||||
self.console.print()
|
||||
self.success(f"File sent: {path.name}")
|
||||
self.active_send = None
|
||||
|
||||
# ── File Transfer: Receiving ─────────────────────────────────────────
|
||||
|
||||
def _handle_file_protocol(self, ft_data: dict, sender: str) -> bool:
|
||||
"""Handle a file transfer protocol message. Returns True if handled."""
|
||||
ft_type = ft_data.get("_ft")
|
||||
transfer_id = ft_data.get("id")
|
||||
|
||||
if ft_type == "offer":
|
||||
if sender == self.username:
|
||||
return True # ignore our own offer echo
|
||||
self.pending_offer = {
|
||||
"id": transfer_id,
|
||||
"name": ft_data.get("name"),
|
||||
"size": ft_data.get("size"),
|
||||
"sha256": ft_data.get("sha256"),
|
||||
"from": sender,
|
||||
}
|
||||
self.transfer_meta[transfer_id] = self.pending_offer
|
||||
self.received_chunks[transfer_id] = []
|
||||
self.render_messages()
|
||||
return True
|
||||
|
||||
elif ft_type == "accept":
|
||||
if self.active_send and self.active_send["id"] == transfer_id:
|
||||
# Someone accepted our offer — start sending chunks
|
||||
asyncio.create_task(
|
||||
self._send_file_chunks(
|
||||
self._ws_ref,
|
||||
self.active_send["id"],
|
||||
self.active_send["path"],
|
||||
)
|
||||
)
|
||||
return True
|
||||
|
||||
elif ft_type == "reject":
|
||||
if self.active_send and self.active_send["id"] == transfer_id:
|
||||
self.error(f"{sender} rejected the file transfer.")
|
||||
self.active_send = None
|
||||
if self.pending_offer and self.pending_offer["id"] == transfer_id:
|
||||
self.pending_offer = None
|
||||
self.render_messages()
|
||||
return True
|
||||
|
||||
elif ft_type == "chunk":
|
||||
if transfer_id in self.received_chunks:
|
||||
chunk_data = base64.b64decode(ft_data.get("data", ""))
|
||||
self.received_chunks[transfer_id].append(chunk_data)
|
||||
meta = self.transfer_meta.get(transfer_id, {})
|
||||
total = meta.get("size", 0)
|
||||
received = sum(len(c) for c in self.received_chunks[transfer_id])
|
||||
pct = int(received * 100 / total) if total else 0
|
||||
self.console.print(
|
||||
f"\r[cyan]Receiving: {pct}% ({_human_size(received)}/{_human_size(total)})[/]",
|
||||
end="",
|
||||
)
|
||||
return True
|
||||
|
||||
elif ft_type == "done":
|
||||
if transfer_id in self.received_chunks:
|
||||
self._finalize_receive(transfer_id)
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def _finalize_receive(self, transfer_id: str) -> None:
|
||||
meta = self.transfer_meta.get(transfer_id, {})
|
||||
chunks = self.received_chunks.pop(transfer_id, [])
|
||||
file_data = b"".join(chunks)
|
||||
|
||||
# Verify integrity
|
||||
actual_sha = hashlib.sha256(file_data).hexdigest()
|
||||
expected_sha = meta.get("sha256", "")
|
||||
if actual_sha != expected_sha:
|
||||
self.console.print()
|
||||
self.error(f"SHA-256 mismatch! File corrupted. Expected {expected_sha[:16]}..., got {actual_sha[:16]}...")
|
||||
return
|
||||
|
||||
# Save file
|
||||
self.download_dir.mkdir(parents=True, exist_ok=True)
|
||||
filename = meta.get("name", f"file_{transfer_id[:8]}")
|
||||
save_path = self.download_dir / filename
|
||||
|
||||
# Avoid overwriting — append number if exists
|
||||
if save_path.exists():
|
||||
stem = save_path.stem
|
||||
suffix = save_path.suffix
|
||||
i = 1
|
||||
while save_path.exists():
|
||||
save_path = self.download_dir / f"{stem}_{i}{suffix}"
|
||||
i += 1
|
||||
|
||||
save_path.write_bytes(file_data)
|
||||
self.console.print()
|
||||
self.success(f"File saved: {save_path} ({_human_size(len(file_data))}) — SHA-256 verified")
|
||||
self.pending_offer = None
|
||||
self.transfer_meta.pop(transfer_id, None)
|
||||
self.render_messages()
|
||||
|
||||
# ── Core Loops ───────────────────────────────────────────────────────
|
||||
|
||||
async def receive_loop(self, ws) -> None:
|
||||
try:
|
||||
|
|
@ -196,6 +392,18 @@ class Client:
|
|||
self.render_messages()
|
||||
elif msg_type == "message":
|
||||
msg_data = self.decrypt_message(data.get("data", {}))
|
||||
text = msg_data.get("text", "")
|
||||
sender = msg_data.get("username", "unknown")
|
||||
|
||||
# Check if this is a file transfer protocol message
|
||||
if text.startswith('{"_ft":'):
|
||||
try:
|
||||
ft_data = json.loads(text)
|
||||
if self._handle_file_protocol(ft_data, sender):
|
||||
continue
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
self.messages.append(msg_data)
|
||||
self.render_messages()
|
||||
elif msg_type == "user_left":
|
||||
|
|
@ -207,6 +415,7 @@ class Client:
|
|||
self.connected = False
|
||||
|
||||
async def input_loop(self, ws) -> None:
|
||||
self._ws_ref = ws
|
||||
loop = asyncio.get_event_loop()
|
||||
while self.running:
|
||||
try:
|
||||
|
|
@ -214,6 +423,52 @@ class Client:
|
|||
if text.lower() in ("q", "quit", "exit"):
|
||||
self.running = False
|
||||
break
|
||||
|
||||
# File transfer commands
|
||||
if text.startswith("/send "):
|
||||
filepath = text[6:].strip()
|
||||
if filepath:
|
||||
await self._send_offer(ws, filepath)
|
||||
else:
|
||||
self.error("Usage: /send <filepath>")
|
||||
continue
|
||||
|
||||
if text.strip() == "/accept":
|
||||
if self.pending_offer:
|
||||
accept_msg = json.dumps({
|
||||
"_ft": "accept",
|
||||
"id": self.pending_offer["id"],
|
||||
})
|
||||
encrypted = self.room_fernet.encrypt(accept_msg.encode()).decode()
|
||||
await ws.send(encrypted)
|
||||
self.info(f"Accepted transfer: {self.pending_offer['name']}")
|
||||
self.pending_offer = None
|
||||
self.render_messages()
|
||||
else:
|
||||
self.error("No pending file offer to accept.")
|
||||
continue
|
||||
|
||||
if text.strip() == "/reject":
|
||||
if self.pending_offer:
|
||||
reject_msg = json.dumps({
|
||||
"_ft": "reject",
|
||||
"id": self.pending_offer["id"],
|
||||
})
|
||||
encrypted = self.room_fernet.encrypt(reject_msg.encode()).decode()
|
||||
await ws.send(encrypted)
|
||||
self.info("Rejected file transfer.")
|
||||
self.received_chunks.pop(self.pending_offer["id"], None)
|
||||
self.transfer_meta.pop(self.pending_offer["id"], None)
|
||||
self.pending_offer = None
|
||||
self.render_messages()
|
||||
else:
|
||||
self.error("No pending file offer to reject.")
|
||||
continue
|
||||
|
||||
if text.strip().startswith("/"):
|
||||
self.error(f"Unknown command: {text.strip().split()[0]}")
|
||||
continue
|
||||
|
||||
if text.strip():
|
||||
encrypted = self.room_fernet.encrypt(text.encode()).decode()
|
||||
await ws.send(encrypted)
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user