#!/usr/bin/python3
"""Thin vendor C ABI acquisition helper; not a verifier or a software NSM."""
import base64
import ctypes
import json
import sys

def decode(value, limit):
    raw = base64.b64decode(value, validate=True)
    if not raw or len(raw) > limit:
        raise ValueError('NSM_INPUT_LIMIT')
    return (ctypes.c_ubyte * len(raw)).from_buffer_copy(raw)

def main():
    wire = sys.stdin.buffer.read(8193)
    if len(wire) > 8192:
        raise ValueError('NSM_REQUEST_LIMIT')
    value = json.loads(wire)
    if set(value) != {'user_data', 'nonce', 'public_key'}:
        raise ValueError('NSM_REQUEST_FIELDS')
    user = decode(value['user_data'], 32)
    if len(user) != 32:
        raise ValueError('NSM_BINDING_HASH')
    nonce = decode(value['nonce'], 128)
    key = decode(value['public_key'], 1024)
    # This path must be measured into the EIF with the Python runtime and helper.
    lib = ctypes.CDLL('/kit/lib/libnsm.so')
    lib.nsm_lib_init.restype = ctypes.c_int32
    lib.nsm_lib_exit.argtypes = [ctypes.c_int32]
    ptr = ctypes.POINTER(ctypes.c_ubyte)
    lib.nsm_get_attestation_doc.argtypes = [ctypes.c_int32, ptr, ctypes.c_uint32,
        ptr, ctypes.c_uint32, ptr, ctypes.c_uint32, ptr, ctypes.POINTER(ctypes.c_uint32)]
    lib.nsm_get_attestation_doc.restype = ctypes.c_int32
    fd = lib.nsm_lib_init()
    if fd < 0:
        raise ValueError('HARDWARE_NSM_REQUIRED')
    try:
        document = (ctypes.c_ubyte * 16384)()
        size = ctypes.c_uint32(len(document))
        status = lib.nsm_get_attestation_doc(fd, user, len(user), nonce, len(nonce),
            key, len(key), document, ctypes.byref(size))
        if status != 0 or not 1 <= size.value <= len(document):
            raise ValueError('NSM_ATTESTATION_FAILED')
        encoded = base64.urlsafe_b64encode(bytes(document[:size.value])).decode().rstrip('=')
        print(json.dumps({'document': encoded}))
    finally:
        lib.nsm_lib_exit(fd)

if __name__ == '__main__':
    try:
        main()
    except Exception:
        print('HARDWARE_NSM_ACQUISITION_FAILED', file=sys.stderr)
        sys.exit(1)
