<?php
/**
 * ID token validation.
 *
 * @package M365_Login
 */

defined( 'ABSPATH' ) || exit;

/**
 * Verifies RS256 signed JWTs against the Microsoft JWKS.
 */
final class M365_Login_JWT {

	const LEEWAY = 120; // Seconds of clock skew tolerated.

	/**
	 * Base64url decode.
	 *
	 * @param string $data Data.
	 * @return string|false
	 */
	public static function b64url_decode( $data ) {
		$data = strtr( $data, '-_', '+/' );
		$pad  = strlen( $data ) % 4;
		if ( $pad ) {
			$data .= str_repeat( '=', 4 - $pad );
		}
		// phpcs:ignore WordPress.PHP.DiscouragedPHPFunctions.obfuscation_base64_decode
		return base64_decode( $data, true );
	}

	/**
	 * Base64url encode.
	 *
	 * @param string $data Data.
	 * @return string
	 */
	public static function b64url_encode( $data ) {
		// phpcs:ignore WordPress.PHP.DiscouragedPHPFunctions.obfuscation_base64_encode
		return rtrim( strtr( base64_encode( $data ), '+/', '-_' ), '=' );
	}

	/**
	 * Decodes and verifies a JWT.
	 *
	 * @param string $jwt       Compact serialised token.
	 * @param array  $jwks      JWKS document (array with 'keys').
	 * @param array  $expected  Expected claims: 'aud', 'nonce', 'tenant' (optional GUID).
	 * @return array|WP_Error Claims on success.
	 */
	public static function verify( $jwt, $jwks, $expected ) {
		$parts = explode( '.', $jwt );
		if ( 3 !== count( $parts ) ) {
			return new WP_Error( 'jwt_format', 'Malformed token.' );
		}

		list( $h64, $p64, $s64 ) = $parts;

		$header    = json_decode( self::b64url_decode( $h64 ), true );
		$claims    = json_decode( self::b64url_decode( $p64 ), true );
		$signature = self::b64url_decode( $s64 );

		if ( ! is_array( $header ) || ! is_array( $claims ) || false === $signature ) {
			return new WP_Error( 'jwt_format', 'Malformed token.' );
		}

		// Only RS256 is accepted; refuse "none" and anything HMAC based.
		if ( empty( $header['alg'] ) || 'RS256' !== $header['alg'] ) {
			return new WP_Error( 'jwt_alg', 'Unsupported token algorithm.' );
		}
		if ( empty( $header['kid'] ) || ! is_string( $header['kid'] ) ) {
			return new WP_Error( 'jwt_kid', 'Token has no key ID.' );
		}

		$public_key = self::find_key( $jwks, $header['kid'] );
		if ( ! $public_key ) {
			return new WP_Error( 'jwt_unknown_kid', 'Signing key not found.' );
		}

		$ok = openssl_verify( $h64 . '.' . $p64, $signature, $public_key, OPENSSL_ALGO_SHA256 );
		if ( 1 !== $ok ) {
			return new WP_Error( 'jwt_signature', 'Token signature is invalid.' );
		}

		$now = time();

		if ( empty( $claims['exp'] ) || ! is_numeric( $claims['exp'] ) || ( (int) $claims['exp'] + self::LEEWAY ) < $now ) {
			return new WP_Error( 'jwt_expired', 'Token has expired.' );
		}
		if ( isset( $claims['nbf'] ) && is_numeric( $claims['nbf'] ) && ( (int) $claims['nbf'] - self::LEEWAY ) > $now ) {
			return new WP_Error( 'jwt_nbf', 'Token is not valid yet.' );
		}
		if ( isset( $claims['iat'] ) && is_numeric( $claims['iat'] ) && ( (int) $claims['iat'] - self::LEEWAY ) > $now ) {
			return new WP_Error( 'jwt_iat', 'Token issued in the future.' );
		}

		// Audience must be our client ID.
		$aud = isset( $claims['aud'] ) ? $claims['aud'] : null;
		if ( is_array( $aud ) ) {
			$aud_ok = in_array( $expected['aud'], $aud, true );
		} else {
			$aud_ok = is_string( $aud ) && hash_equals( $expected['aud'], $aud );
		}
		if ( ! $aud_ok ) {
			return new WP_Error( 'jwt_aud', 'Token audience mismatch.' );
		}

		// Tenant / issuer: v2.0 issuer is https://login.microsoftonline.com/{tid}/v2.0.
		if ( empty( $claims['tid'] ) || ! is_string( $claims['tid'] ) || ! M365_Login_Settings::is_guid( $claims['tid'] ) ) {
			return new WP_Error( 'jwt_tid', 'Token has no tenant ID.' );
		}
		if ( ! empty( $expected['tenant'] ) && ! hash_equals( strtolower( $expected['tenant'] ), strtolower( $claims['tid'] ) ) ) {
			return new WP_Error( 'jwt_tenant', 'Token was issued by a different tenant.' );
		}
		$expected_iss = 'https://login.microsoftonline.com/' . strtolower( $claims['tid'] ) . '/v2.0';
		if ( empty( $claims['iss'] ) || ! is_string( $claims['iss'] ) || ! hash_equals( $expected_iss, strtolower( $claims['iss'] ) ) ) {
			return new WP_Error( 'jwt_iss', 'Token issuer mismatch.' );
		}

		// Nonce binds the token to the login attempt.
		if ( empty( $claims['nonce'] ) || ! is_string( $claims['nonce'] ) || ! hash_equals( $expected['nonce'], $claims['nonce'] ) ) {
			return new WP_Error( 'jwt_nonce', 'Token nonce mismatch.' );
		}

		return $claims;
	}

	/**
	 * Finds a key by kid and returns an OpenSSL public key resource/object.
	 *
	 * @param array  $jwks JWKS document.
	 * @param string $kid  Key ID.
	 * @return mixed|null
	 */
	private static function find_key( $jwks, $kid ) {
		if ( empty( $jwks['keys'] ) || ! is_array( $jwks['keys'] ) ) {
			return null;
		}
		foreach ( $jwks['keys'] as $key ) {
			if ( ! is_array( $key ) || empty( $key['kid'] ) || ! hash_equals( (string) $key['kid'], $kid ) ) {
				continue;
			}
			if ( isset( $key['kty'] ) && 'RSA' !== $key['kty'] ) {
				continue;
			}
			if ( isset( $key['use'] ) && 'sig' !== $key['use'] ) {
				continue;
			}

			// Prefer the embedded certificate, fall back to modulus/exponent.
			if ( ! empty( $key['x5c'][0] ) && is_string( $key['x5c'][0] ) ) {
				$pem = "-----BEGIN CERTIFICATE-----\n" . chunk_split( $key['x5c'][0], 64, "\n" ) . "-----END CERTIFICATE-----\n";
				$pub = openssl_pkey_get_public( $pem );
				if ( $pub ) {
					return $pub;
				}
			}
			if ( ! empty( $key['n'] ) && ! empty( $key['e'] ) ) {
				$pem = self::rsa_pem_from_components( $key['n'], $key['e'] );
				if ( $pem ) {
					$pub = openssl_pkey_get_public( $pem );
					if ( $pub ) {
						return $pub;
					}
				}
			}
		}
		return null;
	}

	/**
	 * Builds a PEM encoded SubjectPublicKeyInfo from JWK modulus / exponent.
	 *
	 * @param string $n Base64url modulus.
	 * @param string $e Base64url exponent.
	 * @return string|false
	 */
	private static function rsa_pem_from_components( $n, $e ) {
		$modulus  = self::b64url_decode( $n );
		$exponent = self::b64url_decode( $e );
		if ( false === $modulus || false === $exponent ) {
			return false;
		}

		$modulus  = self::der_integer( $modulus );
		$exponent = self::der_integer( $exponent );

		$rsa_key = self::der_seq( $modulus . $exponent );
		// rsaEncryption OID 1.2.840.113549.1.1.1 + NULL.
		$alg_id  = self::der_seq( "\x06\x09\x2a\x86\x48\x86\xf7\x0d\x01\x01\x01\x05\x00" );
		$bit_str = "\x03" . self::der_len( strlen( $rsa_key ) + 1 ) . "\x00" . $rsa_key;
		$spki    = self::der_seq( $alg_id . $bit_str );

		// phpcs:ignore WordPress.PHP.DiscouragedPHPFunctions.obfuscation_base64_encode
		return "-----BEGIN PUBLIC KEY-----\n" . chunk_split( base64_encode( $spki ), 64, "\n" ) . "-----END PUBLIC KEY-----\n";
	}

	/**
	 * DER length encoding.
	 *
	 * @param int $len Length.
	 * @return string
	 */
	private static function der_len( $len ) {
		if ( $len < 128 ) {
			return chr( $len );
		}
		$bytes = ltrim( pack( 'N', $len ), "\x00" );
		return chr( 0x80 | strlen( $bytes ) ) . $bytes;
	}

	/**
	 * DER SEQUENCE.
	 *
	 * @param string $content Content.
	 * @return string
	 */
	private static function der_seq( $content ) {
		return "\x30" . self::der_len( strlen( $content ) ) . $content;
	}

	/**
	 * DER INTEGER (unsigned big-endian input).
	 *
	 * @param string $bytes Raw bytes.
	 * @return string
	 */
	private static function der_integer( $bytes ) {
		$bytes = ltrim( $bytes, "\x00" );
		if ( '' === $bytes || ( ord( $bytes[0] ) & 0x80 ) ) {
			$bytes = "\x00" . $bytes;
		}
		return "\x02" . self::der_len( strlen( $bytes ) ) . $bytes;
	}
}