#!/usr/bin/env python3
"""Chappa Facturare — sign a set of claims as a compact JWS (ES256), with the
public key in the header.

    pip install cryptography
    python chappa_sign.py <private-key.pem> '<claims JSON>' [lifetime-seconds]

Adds iat, exp (default +300 s, max 900) and a random jti when missing.
"""
import base64
import json
import sys
import time
import uuid

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.asymmetric.utils import decode_dss_signature


def b64u(b: bytes) -> str:
    return base64.urlsafe_b64encode(b).rstrip(b"=").decode()


def main() -> None:
    if len(sys.argv) < 3:
        print("usage: python chappa_sign.py <private-key.pem> '<claims JSON>' [lifetime-seconds]", file=sys.stderr)
        sys.exit(2)
    key = serialization.load_pem_private_key(open(sys.argv[1], "rb").read(), None)
    if not isinstance(key, ec.EllipticCurvePrivateKey) or key.curve.name != "secp256r1":
        print("the key must be EC P-256", file=sys.stderr)
        sys.exit(2)
    claims = json.loads(sys.argv[2])
    lifetime = min(int(sys.argv[3]) if len(sys.argv) > 3 else 300, 900)
    now = int(time.time())
    claims.setdefault("iat", now)
    claims.setdefault("exp", now + lifetime)
    claims.setdefault("jti", str(uuid.uuid4()))
    nums = key.public_key().public_numbers()
    jwk = {"kty": "EC", "crv": "P-256", "x": b64u(nums.x.to_bytes(32, "big")), "y": b64u(nums.y.to_bytes(32, "big"))}
    header = {"alg": "ES256", "typ": "JWT", "jwk": jwk}
    signing_input = b64u(json.dumps(header, separators=(",", ":")).encode()) + "." + \
        b64u(json.dumps(claims, separators=(",", ":")).encode())
    r, s = decode_dss_signature(key.sign(signing_input.encode(), ec.ECDSA(hashes.SHA256())))
    print(signing_input + "." + b64u(r.to_bytes(32, "big") + s.to_bytes(32, "big")))


if __name__ == "__main__":
    main()