"""Client: send a file to the relay, or receive it from the relay.

  python client.py send file.zip --relay http://HOST:8888 --token T --secret PASS
      -> prints a transfer ID
  python client.py receive TRANSFER_ID --relay http://HOST:8888 --token T --secret PASS --out ./dl

Files are split into chunks and encrypted on your machine (AES via Fernet,
key derived from --secret). The relay never sees plaintext.
Requires:  pip install cryptography
"""
import argparse
import base64
import hashlib
import json
import secrets
import sys
import urllib.error
import urllib.request
from pathlib import Path

from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

CHUNK = 8 * 1024 * 1024


def fernet(secret: str, salt: bytes) -> Fernet:
    kdf = PBKDF2HMAC(hashes.SHA256(), 32, salt, 300_000)
    return Fernet(base64.urlsafe_b64encode(kdf.derive(secret.encode())))


def call(method, url, token, data=None):
    req = urllib.request.Request(url, data=data, method=method,
                                 headers={"Authorization": f"Bearer {token}"})
    for attempt in range(5):
        try:
            with urllib.request.urlopen(req, timeout=120) as r:
                return r.read()
        except urllib.error.HTTPError as e:
            if e.code < 500:
                raise
        except urllib.error.URLError:
            pass
    raise RuntimeError(f"failed: {method} {url}")


def send(a):
    path = Path(a.file)
    tid = secrets.token_urlsafe(12)
    salt = secrets.token_bytes(16)
    f = fernet(a.secret, salt)
    size = path.stat().st_size
    n = max(1, -(-size // CHUNK))
    sha = hashlib.sha256()
    base = f"{a.relay.rstrip('/')}/transfers/{tid}"
    with open(path, "rb") as fh:
        for i in range(n):
            part = fh.read(CHUNK)
            sha.update(part)
            call("PUT", f"{base}/chunks/{i}", a.token, f.encrypt(part))
            print(f"\rchunk {i + 1}/{n}", end="", flush=True)
    meta = {"name": path.name, "size": size, "chunks": n, "sha256": sha.hexdigest(),
            "salt": base64.b64encode(salt).decode()}
    call("PUT", f"{base}/meta", a.token, json.dumps(meta).encode())
    print(f"\nDone. Transfer ID: {tid}")


def receive(a):
    base = f"{a.relay.rstrip('/')}/transfers/{a.id}"
    info = json.loads(call("GET", base, a.token))
    if "chunks" not in info:
        sys.exit("Upload not finished yet (no metadata).")
    f = fernet(a.secret, base64.b64decode(info["salt"]))
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    target = out / Path(info["name"]).name
    sha = hashlib.sha256()
    with open(target, "wb") as fh:
        for i in range(info["chunks"]):
            try:
                part = f.decrypt(call("GET", f"{base}/chunks/{i}", a.token))
            except InvalidToken:
                sys.exit("Wrong secret or corrupted data.")
            sha.update(part)
            fh.write(part)
            print(f"\rchunk {i + 1}/{info['chunks']}", end="", flush=True)
    if sha.hexdigest() != info["sha256"]:
        sys.exit("\nChecksum mismatch!")
    print(f"\nSaved to {target} (checksum OK)")
    if a.delete:
        call("DELETE", base, a.token)


p = argparse.ArgumentParser()
sub = p.add_subparsers(dest="cmd", required=True)
for name in ("send", "receive"):
    s = sub.add_parser(name)
    s.add_argument("file" if name == "send" else "id")
    s.add_argument("--relay", required=True)
    s.add_argument("--token", required=True)
    s.add_argument("--secret", required=True)
    if name == "receive":
        s.add_argument("--out", default=".")
        s.add_argument("--delete", action="store_true", help="remove from relay after download")
a = p.parse_args()
send(a) if a.cmd == "send" else receive(a)