Keyless agent -> Web SDK tutorial

This page is a tutorial for integrators who want to enroll user selfies captured outside of Keyless using the Keyless agent and then allow those same users to authenticate on their web app/SDK

"""
Tutorial to support integrators in achieving cooperation between the Keyless Agent and the Authentication Service.
It is intended to be used as executable tutorial code, which can be run in a Python environment.
Best way is to read it from top to bottom as concepts are introduced in a sequence.
If you have a tool like uv, hatch, etc. that can read inline metadata then you can execute this script directly.
Otherwise, install the required dependencies.

Enjoy the tutorial!
"""
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "cryptography",
#     "httpx",
# ]
# ///

import argparse
import logging
import os
from pathlib import Path
from pprint import pprint
from typing import Literal, TypedDict

import httpx
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.ciphers.aead import AESGCMSIV as AES_GCM_SIV
from cryptography.hazmat.primitives.serialization import load_pem_public_key

logging.basicConfig(level=logging.INFO)

# This tutorial demonstrates the cooperation between the Keyless Agent and the Authentication Service.

# The process is fairly simple:
# 1. Use Keyless Agent to create a user get client state
# 2. Import the client state to the Authentication Service
# 3. Profit.


#
# Keyless Agent
#

# Keyless agent is a http server that gets the user's face as an image creates
# and returns a Keyless user.
#
# There are two flavors of this functionality, online and offline.
# This tutorial will cover the offline version.

# To start we need to have user image in a JPEG format (other formats are supported).
# And we need to have a scenario in mind that describes how the image was taken.
# Each scenario has different requirements and is used to determine the level of trust in the image.
#
# * `SELFIE` - the image was taken by the user themselves, usually just moments ago.
# * `TRUSTED_SOURCE` - the image was already validated by trusted party.
# * `DOCUMENT` - the image is taken/copied from a document. Like passport, government ID, etc.

type Scenario = Literal["SELFIE", "TRUSTED_SOURCE", "DOCUMENT"]

# After that its a simple HTTP POST request to the Keyless Agent.
# The response is JSON with the user's ID and other information needed to continue.


def create_user_offline(
    keyless_agent_url: str, scenario: Scenario, image_jpeg: bytes
) -> "OfflineEnrollmentUser":
    """Create a user on the Keyless Agent using offline flow."""
    result = httpx.post(
        f"{keyless_agent_url}/v1/offline-enrollment",
        headers={"Scenario": scenario, "Content-Type": "image/jpeg"},
        content=image_jpeg,
    )
    assert result.status_code == 200, result.text
    return result.json()


# During this call Keyless Agent process the image and generate two parts of a user
# client state and sever state. As this is the offline flow, the user is not really
# registered in the system yet and both parts are needed to continue.


class OfflineEnrollmentUser(TypedDict):
    """Structure of the response from the Keyless Agent."""

    keylessId: str
    """The id of the user in the Keyless system."""

    clientState: str
    """Opaque string used to initialize the SDK."""

    serverState: str
    """Opaque string used to register the user."""


#
# Register the server state and activate user
#

# As mentioned before, the user is not really registered in the system yet.
# To do that we need to send the server state to the Operations Service.
#
# Operation service than register the user to be used.
#
# Api calls to operations service is authorized using secret key, provided by ke keyless.
#
# Note: if the Keyless Agent is configured for online flow, than you can use
# `/v1/online-enrollment` enrollment endpoint. The user is registered
# automatically and only client state is returned.


def register_the_user(
    operation_service_url: str, secret_key: str, user: OfflineEnrollmentUser
) -> None:
    """Register the user on the Operation Service."""
    #
    # First we need to send the server state to the operation service.
    response = httpx.post(
        f"{operation_service_url}/v2/uncommitted-users",
        content=user["serverState"].encode("utf-8"),
        headers={"Content-Type": "application/json", "X-Api-Key": secret_key},
    )
    assert response.status_code == 201, response.text

    # Than we need to activate the user.
    response = httpx.post(
        f"{operation_service_url}/v2/uncommitted-users/{user['keylessId']}/commit",
        headers={"X-Api-Key": secret_key},
    )
    assert response.status_code == 204, response.text


#
# Add a user to the Authentication Service
#

# Now the client state can be used in any SDK. This tutorial focusses on the
# Web SDK. For that we need to import the client state to the Authentication Service.
#
# The client state contains sensitive information and must be encrypted before sending.
# The encryption is done using the public key provided by keyless.

# The client state encryption is done in two steps.
#
# 1. The client state is encrypted using random AES key.
# 2. The AES key is encrypted using the public key provided by Keyless.

# In this tutorial we are using recommended AES-GCM-SIV encryption algorithm.


# first generate random AES key, 128 bits is the recommended size
def generate_aes_key() -> bytes:
    return AES_GCM_SIV.generate_key(128)


# than encrypt the AES key using the public key
def encrypt_client_state(plain_key: bytes, user: OfflineEnrollmentUser) -> bytes:
    # than generate random nonce 12 bytes long
    nonce = os.urandom(12)

    # encrypt the client state using the AES key and nonce
    encrypted_client_state = AES_GCM_SIV(plain_key).encrypt(
        nonce, user["clientState"].encode("utf-8"), None
    )

    # the nonce is prepended to the encrypted data
    return nonce + encrypted_client_state


# Now we need to encrypt the AES key using the public key provided by Keyless.
#
# For demos and tutorials, the public key is available at the endpoint `/v1/customers/{customer}/client-state-encryption`.
# But it should not be used in production to prevent attackers to modify it.
#
# customer is the name of the customer/integrator in the Keyless system.


def get_client_state_encryption(
    authentication_service_url: str, secret_key: str, customer: str
) -> "ClientStateEncryption":
    response = httpx.get(
        f"{authentication_service_url}/v1/customers/{customer}/client-state-encryption",
        headers={"kl-api-key": secret_key},
    )
    assert response.status_code == 200, response.text
    return response.json()


class ClientStateEncryption(TypedDict):
    """Structure of the response from the Authentication Service."""

    keyId: str
    """The id of the key used to encrypt the client state."""
    supportedAlgorithms: list[str]
    """List of supported encryption algorithms."""
    publicKey: str
    """The public key used to encrypt the client state."""


# Now we can encrypt the AES key using the public key.
#
# For this we are using the RSAES-OAEP-SHA-256 algorithm.


def encrypt_aes_key(
    plain_key: bytes, client_state_encryption: ClientStateEncryption
) -> bytes:
    key_algorithm = client_state_encryption["supportedAlgorithms"][0]

    # just a sanity check that the algorithm is what we expect
    assert key_algorithm == "RSAES-OAEP-SHA-256", key_algorithm

    # load the public key
    pk = load_pem_public_key(client_state_encryption["publicKey"].encode("ascii"))

    # encrypt the AES key using the public key and the RSAES-OAEP-SHA-256 algorithm
    return pk.encrypt(
        plain_key,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None,
        ),
    )


# Now we finally have all the pieces to import the client state to the Authentication Service.
#
# Unlike the other SDKs the Web SDK operates on the usernames and not the user ids (keylessId).


def import_user(
    authentication_service_url: str,
    secret_key: str,
    customer: str,
    username: str,
    client_state_encryption: ClientStateEncryption,
    user: OfflineEnrollmentUser,
):
    plain_key = generate_aes_key()
    client_state = encrypt_client_state(plain_key, user)
    encrypted_key = encrypt_aes_key(plain_key, client_state_encryption)

    response = httpx.post(
        f"{authentication_service_url}/v1/users/{customer}/{username}/import-client-state",
        headers={
            "kl-api-key": secret_key,
            "content-type": "application/octet-stream",
            # the headers are used to notify the Authentication Service about the encryption
            # used
            "Kl-Key-Id": client_state_encryption["keyId"],
            "Kl-Key-Algorithm": "RSAES-OAEP-SHA-256",
            "Kl-Client-State-Key": encrypted_key.hex(),
            "Kl-Client-State-Algorithm": "AES-GCM-SIV",
        },
        content=client_state,
    )
    assert response.status_code == 200, response.text

    return response.json()


# And that is it. The user is now registered in the Authentication Service and can be used in the Web SDK.
#
# If you wan to try it run this script.


def main():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--keyless-agent-url", required=True, help="URL of the Keyless Agent"
    )
    parser.add_argument(
        "--operation-service-url", required=True, help="URL of the Operation Service"
    )
    parser.add_argument(
        "--authentication-service-url",
        required=True,
        help="URL of the Authentication Service",
    )
    parser.add_argument(
        "--secret-key", required=True, help="Secret key for the API calls"
    )
    parser.add_argument("--image-file", required=True)

    parser.add_argument("customer", help="Customer name")
    parser.add_argument("username", help="Username of the user")

    params = parser.parse_args()

    # now the full flow
    user = create_user_offline(
        params.keyless_agent_url, "SELFIE", Path(params.image_file).read_bytes()
    )

    register_the_user(params.operation_service_url, params.secret_key, user)

    response = import_user(
        params.authentication_service_url,
        params.secret_key,
        params.customer,
        params.username,
        get_client_state_encryption(
            params.authentication_service_url, params.secret_key, params.customer
        ),  # note: this is just for demo, encryption configuration should be part of the application configuration
        user,
    )
    pprint(response)


if __name__ == "__main__":
    main()