diff --git a/cmd_chat/client/client.py b/cmd_chat/client/client.py index b1b9214..04c4cc2 100644 --- a/cmd_chat/client/client.py +++ b/cmd_chat/client/client.py @@ -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 | /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 ") + 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)