# receiver.py - nhan file qua TCP. Dung chung cho bai 9. Chay: python3 receiver.py [--tls] [--host IP]
# Co 5 lop bao ve: token, gioi han kich thuoc, lam sach ten file, kiem tra SHA-256, tuy chon TLS.
import argparse
import hashlib
import hmac
import json
import os
import socket
import ssl
from pathlib import Path

MAX_HEADER = 4096
MAX_SIZE = 50 * 1024 * 1024
SAVE_DIR = Path("received")


def read_header(conn):
    data = b""
    while b"\n" not in data:
        chunk = conn.recv(1)
        if not chunk or len(data) > MAX_HEADER:
            raise ValueError("header sai hoac qua dai")
        data += chunk
    return json.loads(data.decode("utf-8"))


def read_exact(conn, size):
    parts, left = [], size
    while left > 0:
        chunk = conn.recv(min(65536, left))
        if not chunk:
            raise ValueError("ket noi dong giua chung")
        parts.append(chunk)
        left -= len(chunk)
    return b"".join(parts)


def handle(conn, token):
    header = read_header(conn)
    if not hmac.compare_digest(str(header.get("token", "")), token):
        raise PermissionError("sai token")
    size = int(header["size"])
    if not 0 <= size <= MAX_SIZE:
        raise ValueError("kich thuoc khong hop le")
    name = os.path.basename(str(header["name"]).replace("\\", "/"))
    if not name or name.startswith("."):
        raise ValueError("ten file khong hop le")
    body = read_exact(conn, size)
    if hashlib.sha256(body).hexdigest() != header["sha256"]:
        raise ValueError("SHA-256 khong khop: du lieu bi loi hoac bi sua")
    SAVE_DIR.mkdir(exist_ok=True)
    (SAVE_DIR / name).write_bytes(body)
    return name, size


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="127.0.0.1", help="IP de nghe; 0.0.0.0 la moi card mang")
    ap.add_argument("--port", type=int, default=9000)
    ap.add_argument("--token", default="doi-token-nay", help="chuoi bi mat chung giua hai ben")
    ap.add_argument("--tls", action="store_true", help="bat TLS, can cert.pem va key.pem")
    args = ap.parse_args()
    context = None
    if args.tls:
        context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        context.load_cert_chain("cert.pem", "key.pem")
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server:
        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server.bind((args.host, args.port))
        server.listen(5)
        print(f"Dang nghe {args.host}:{args.port} ({'TLS' if args.tls else 'KHONG ma hoa'}). Ctrl+C de dung.")
        while True:
            conn, addr = server.accept()
            conn.settimeout(10)
            try:
                with (context.wrap_socket(conn, server_side=True) if context else conn) as c:
                    try:
                        name, size = handle(c, args.token)
                        c.sendall(b"OK\n")
                        print(f"Nhan {name} ({size} byte) tu {addr[0]}")
                    except (ValueError, PermissionError, KeyError, json.JSONDecodeError) as err:
                        c.sendall(f"LOI {err}\n".encode())
                        print(f"Tu choi {addr[0]}: {err}")
            except (OSError, ssl.SSLError) as err:
                print(f"Loi ket noi tu {addr[0]}: {err}")


if __name__ == "__main__":
    main()