<?php
/**
 * Client authentication at the token and revocation endpoints:
 * private_key_jwt (RFC 7523 §2.2; spec 4.1).
 *
 * The client assertion is a short-lived JWT signed with a key from the
 * agent's jwks_uri, with iss = sub = client_id and aud = this store's OAuth
 * issuer (`auth.issuer`) as a single string, at every endpoint that takes it
 * (spec 4.1, after draft-ietf-oauth-rfc7523bis). Its jti is single-use. The
 * agent behind it must resolve (metadata + keys) and must not be blocked.
 *
 * @package AgentAdapt
 */

defined( 'ABSPATH' ) || exit;

/**
 * Client authentication.
 */
final class AgentAdapt_Client_Auth {

	public const ASSERTION_TYPE = 'urn:ietf:params:oauth:client-assertion-type:jwt-bearer';

	/**
	 * Longest client/session assertion lifetime accepted (exp - iat).
	 */
	public const MAX_ASSERTION_LIFETIME = 300;

	/**
	 * Authenticate the client of a token-endpoint style request.
	 *
	 * @param array<string,string> $form Form parameters.
	 * @return array<string,mixed>|WP_Error Resolved agent, or an error whose code is the OAuth error.
	 */
	public static function authenticate( array $form ) {
		$client_id = (string) ( $form['client_id'] ?? '' );
		if ( '' !== $client_id && ! isset( $form['client_assertion_type'] ) && ! isset( $form['client_assertion'] ) && AgentAdapt_Agents::public_clients_allowed() ) {
			// A public client (MCP's own sign-in): nothing to authenticate. The
			// token endpoint limits it to PKCE-protected codes and its refresh
			// token, for Bearer resources only.
			$agent = AgentAdapt_Agents::resolve( $client_id );
			if ( is_wp_error( $agent ) ) {
				return 'rate_limited' === $agent->get_error_code() ? $agent : new WP_Error( 'invalid_client', $agent->get_error_message() );
			}
			return empty( $agent['public'] ) ? new WP_Error( 'invalid_client', 'this client must authenticate with private_key_jwt' ) : $agent;
		}
		if ( '' === $client_id || self::ASSERTION_TYPE !== ( $form['client_assertion_type'] ?? '' ) || empty( $form['client_assertion'] ) ) {
			return new WP_Error( 'invalid_client', 'private_key_jwt client authentication is required' );
		}
		$agent = AgentAdapt_Agents::resolve( $client_id );
		if ( is_wp_error( $agent ) ) {
			return 'rate_limited' === $agent->get_error_code() ? $agent : new WP_Error( 'invalid_client', $agent->get_error_message() );
		}
		$claims = self::verify_agent_jwt( $agent, (string) $form['client_assertion'], 'client' );
		if ( is_wp_error( $claims ) ) {
			return new WP_Error( 'invalid_client', 'client assertion: ' . $claims->get_error_message() );
		}
		if ( ( $claims['sub'] ?? null ) !== $client_id ) {
			return new WP_Error( 'invalid_client', 'client assertion sub must be the client_id' );
		}
		// Spec 4.1: aud is the issuer, as a single string, and nothing else,
		// not even the token endpoint URL RFC 7523 allowed. Nobody verifies
		// the endpoints in an issuer's metadata, so an assertion addressed to
		// an endpoint could be collected by another Company that names ours
		// and replayed here (spec 9.4).
		if ( ! self::aud_is_issuer( $claims ) ) {
			return new WP_Error( 'invalid_client', 'client assertion aud must be this store\'s issuer (auth.issuer), as a single string' );
		}
		return $agent;
	}

	/**
	 * Whether an assertion's aud is exactly this store's issuer identifier,
	 * as one string (spec 4.1 client assertions, 4.2 Session assertions). An
	 * array, an endpoint URL or any other value is refused.
	 *
	 * @param array<string,mixed> $claims Verified claims.
	 */
	public static function aud_is_issuer( array $claims ): bool {
		return is_string( $claims['aud'] ?? null ) && AgentAdapt_Urls::issuer() === $claims['aud'];
	}

	/**
	 * Verify a JWT the agent signed with a key from its jwks_uri, and the
	 * claims every such JWT shares: iss = client_id, iat/exp present and
	 * short-lived, a single-use jti.
	 *
	 * @param array<string,mixed> $agent Resolved agent.
	 * @param string              $jwt   Compact JWS.
	 * @param string              $kind  Replay namespace: client|assertion|browser.
	 * @param string|null         $typ   Required typ header, or null.
	 * @return array<string,mixed>|WP_Error Claims.
	 */
	public static function verify_agent_jwt( array $agent, string $jwt, string $kind, ?string $typ = null ) {
		$peek = AgentAdapt_Jose::peek( $jwt );
		if ( null === $peek ) {
			return new WP_Error( 'malformed', 'not a compact JWS' );
		}
		// Spec §5: a browser assertion is never accepted at the token endpoint,
		// and a session assertion is never accepted as a browser assertion. The
		// browser endpoint asks for its typ explicitly; everyone else refuses it.
		if ( null === $typ && 0 === strcasecmp( (string) ( $peek[0]['typ'] ?? '' ), AgentAdapt_Browser_Session::TYP ) ) {
			return new WP_Error( 'typ', 'a browser assertion is not accepted here' );
		}
		$kid  = isset( $peek[0]['kid'] ) ? (string) $peek[0]['kid'] : null;
		$keys = AgentAdapt_Agents::keys_for( $agent, $kid );
		$v    = AgentAdapt_Jose::verify( $jwt, $keys, $typ );
		if ( is_wp_error( $v ) ) {
			return $v;
		}
		$c = $v['claims'];
		if ( ( $c['iss'] ?? null ) !== $agent['client_id'] ) {
			return new WP_Error( 'iss', 'iss must be the client_id' );
		}
		if ( ! is_int( $c['exp'] ?? null ) || ! is_int( $c['iat'] ?? null ) ) {
			return new WP_Error( 'lifetime', 'iat and exp are required' );
		}
		if ( $c['exp'] - $c['iat'] > self::MAX_ASSERTION_LIFETIME || $c['exp'] <= $c['iat'] ) {
			return new WP_Error( 'lifetime', 'assertion lifetime is too long' );
		}
		$jti = $c['jti'] ?? null;
		if ( ! is_string( $jti ) || strlen( $jti ) < 16 || strlen( $jti ) > 256 ) {
			return new WP_Error( 'jti', 'a jti of at least 128 bits is required' );
		}
		if ( ! AgentAdapt_Replay::claim( $kind, $agent['client_id'], $jti, (int) $c['exp'] ) ) {
			return new WP_Error( 'replay', 'this assertion was already used' );
		}
		return $c;
	}
}