#!/usr/bin/env python3
"""Create a public PXA Store key-registration proof locally; performs no network IO."""
import argparse
import base64
import getpass
import json
import re
from pathlib import Path
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 create_proof(key, challenge):
    if not isinstance(key, ec.EllipticCurvePrivateKey) or not isinstance(key.curve, ec.SECP256R1):
        raise ValueError('A P-256 private key is required')
    if not re.fullmatch(r'[A-Za-z0-9_-]{43}', challenge):
        raise ValueError('Expected the 43-character challenge returned by PXA Store')
    message = b'PXA-STORE-PUBLISHER-KEY\0' + challenge.encode('ascii')
    r, s = decode_dss_signature(key.sign(message, ec.ECDSA(hashes.SHA256())))
    spki = key.public_key().public_bytes(serialization.Encoding.DER, serialization.PublicFormat.SubjectPublicKeyInfo)
    return {'challenge': challenge, 'publisher_spki': base64.b64encode(spki).decode(),
            'signature': base64.b64encode(r.to_bytes(32, 'big') + s.to_bytes(32, 'big')).decode()}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('key', type=Path, help='Local P-256 PEM private key; never uploaded')
    args = parser.parse_args()
    pem = args.key.read_bytes()
    password = getpass.getpass('PEM password: ').encode() if b'ENCRYPTED' in pem else None
    key = serialization.load_pem_private_key(pem, password=password)
    challenge = input('Paste challenge: ').strip()
    print(json.dumps(create_proof(key, challenge), ensure_ascii=False))


if __name__ == '__main__':
    main()