#!/usr/bin/env python3
"""
FZP1 client for the Finzep payout API - Python 3.8+  (pip install cryptography requests)
Specification: payout-api-guide.html (section "Encryption")

Save this file as fzp_client.py. Use as a library:

    from fzp_client import FzpClient
    client = FzpClient(base_url="https://api.example.com", api_key="fzp_...", client_key_id="ck_...",
                       client_private_pem=open("client-private.pem").read(), server_key_id="sk_...", server_public_pem="...")
    status, answer = client.post("/API/Payout", {"APIRequestID": "REQ-1", "TransMode": "IMPS", ...})

or from the command line:
    python fzp_client.py keygen <dir>
    python fzp_client.py build-request --client-key-id ck_.. --server-key-id sk_.. --path /API/Payout \
           --client-private client-private.pem --server-public server-public.pem --json '{"APIRequestID":"R1",...}'
    python fzp_client.py open-response --request-envelope request.json --response response.json \
           --client-private client-private.pem --server-public server-public.pem
    python fzp_client.py open-webhook --body webhook-body.json --event-id <X-Webhook-Id header> \
           --client-private client-private.pem --server-public server-public.pem
"""
import base64
import datetime
import json
import os
import sys

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, rsa
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

VERSION = "FZP1"
PSS = padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=32)
OAEP = padding.OAEP(mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None)


def _b64(data: bytes) -> str:
    return base64.b64encode(data).decode("ascii")


def signed_text(kind: str, path_or_request_nonce: str, e: dict) -> bytes:
    parts = [VERSION, kind, path_or_request_nonce, e["clientKeyId"], e["serverKeyId"], e["timestamp"],
             e["nonce"], e["encryptedKey"], e["iv"], e["payload"]]
    return "\n".join(parts).encode("utf-8")


def _aad(e: dict) -> bytes:
    return f'{e["timestamp"]}\n{e["nonce"]}'.encode("utf-8")


def build_request(client_key_id, server_key_id, path, json_text, client_private_pem, server_public_pem) -> dict:
    """Encrypts `json_text` for the server and signs it with the client's private key. Returns the envelope (a dict)."""
    aes_key, iv = os.urandom(32), os.urandom(12)
    e = {
        "clientKeyId": client_key_id,
        "serverKeyId": server_key_id,
        "timestamp": datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "nonce": _b64(os.urandom(16)),
        "iv": _b64(iv),
    }
    server_public = serialization.load_pem_public_key(server_public_pem.encode())
    e["encryptedKey"] = _b64(server_public.encrypt(aes_key, OAEP))
    e["payload"] = _b64(AESGCM(aes_key).encrypt(iv, json_text.encode("utf-8"), _aad(e)))  # ciphertext || 16-byte tag
    client_private = serialization.load_pem_private_key(client_private_pem.encode(), password=None)
    e["signature"] = _b64(client_private.sign(signed_text("REQUEST", path, e), PSS, hashes.SHA256()))
    return e


def open_response(response_body, request_envelope, client_private_pem, server_public_pem) -> str:
    """Verifies the server's signature, then decrypts its answer. Raises if anything does not check out."""
    return _open("RESPONSE", request_envelope["nonce"], response_body, client_private_pem, server_public_pem)


def open_webhook(body, event_id, client_private_pem, server_public_pem) -> str:
    """Verifies and decrypts a WEBHOOK we sent you. event_id is the X-Webhook-Id header of the call. Returns the JSON text."""
    return _open("WEBHOOK", event_id, body, client_private_pem, server_public_pem)


def _open(kind, third, body, client_private_pem, server_public_pem) -> str:
    r = json.loads(body) if isinstance(body, str) else body
    server_public = serialization.load_pem_public_key(server_public_pem.encode())
    try:
        server_public.verify(base64.b64decode(r["signature"]), signed_text(kind, third, r), PSS, hashes.SHA256())
    except InvalidSignature:
        raise ValueError("Server signature does not verify - do not trust this response")
    client_private = serialization.load_pem_private_key(client_private_pem.encode(), password=None)
    aes_key = client_private.decrypt(base64.b64decode(r["encryptedKey"]), OAEP)
    return AESGCM(aes_key).decrypt(base64.b64decode(r["iv"]), base64.b64decode(r["payload"]), _aad(r)).decode("utf-8")


class FzpClient:
    def __init__(self, base_url, api_key, client_key_id, client_private_pem, server_key_id, server_public_pem):
        self.base_url, self.api_key = base_url.rstrip("/"), api_key
        self.client_key_id, self.client_private_pem = client_key_id, client_private_pem
        self.server_key_id, self.server_public_pem = server_key_id, server_public_pem

    def post(self, path, body):
        """Sends one encrypted call. Returns (http_status, answer_dict)."""
        import requests  # pip install requests
        req = build_request(self.client_key_id, self.server_key_id, path, json.dumps(body), self.client_private_pem, self.server_public_pem)
        resp = requests.post(self.base_url + path, data=json.dumps(req),
                             headers={"Content-Type": "application/json", "X-Api-Key": self.api_key, "X-Encrypted": "1"}, timeout=60)
        if "X-Encrypted" in resp.headers:       # errors found before the message could be trusted come back as plain JSON
            return resp.status_code, json.loads(open_response(resp.text, req, self.client_private_pem, self.server_public_pem))
        return resp.status_code, resp.json()


def keygen():
    key = rsa.generate_private_key(public_exponent=65537, key_size=3072)
    private_pem = key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()).decode()
    public_pem = key.public_key().public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo).decode()
    return public_pem, private_pem


if __name__ == "__main__":
    cmd, rest = (sys.argv[1] if len(sys.argv) > 1 else ""), sys.argv[2:]

    def arg(name):
        return rest[rest.index("--" + name) + 1] if ("--" + name) in rest else None

    def read(name):
        with open(arg(name), encoding="utf-8") as f:
            return f.read()

    if cmd == "keygen":
        out = rest[0] if rest else "."
        pub, priv = keygen()
        with open(os.path.join(out, "client-private.pem"), "w") as f:
            f.write(priv)
        with open(os.path.join(out, "client-public.pem"), "w") as f:
            f.write(pub)
        print("Wrote client-private.pem (KEEP SECRET) and client-public.pem (register this one) in " + out)
    elif cmd == "build-request":
        sys.stdout.write(json.dumps(build_request(arg("client-key-id"), arg("server-key-id"), arg("path"), arg("json"),
                                                  read("client-private"), read("server-public"))))
    elif cmd == "open-response":
        sys.stdout.write(open_response(read("response"), json.loads(read("request-envelope")), read("client-private"), read("server-public")))
    elif cmd == "open-webhook":
        sys.stdout.write(open_webhook(read("body"), arg("event-id"), read("client-private"), read("server-public")))
    else:
        sys.exit("usage: fzp_client.py keygen <dir> | build-request ... | open-response ... | open-webhook ...")