Code Coverage
 
Lines
Functions and Methods
Classes and Traits
Total
100.00%
46 / 46
100.00%
5 / 5
 CRAP
100.00%
1 / 1
WP_REST_Jetpack_AI_JWT
100.00%
46 / 46
100.00%
5 / 5
 9
100.00%
1 / 1
 __construct
100.00%
2 / 2
100.00%
1 / 1
 1
 register_rest_route
100.00%
13 / 13
100.00%
1 / 1
 2
 is_rest_endpoint_registered
100.00%
4 / 4
100.00%
1 / 1
 1
 permission_callback
100.00%
1 / 1
100.00%
1 / 1
 2
 get_jwt
100.00%
26 / 26
100.00%
1 / 1
 3
1<?php
2/**
3 * WP_REST_Jetpack_AI_JWT file.
4 *
5 * @package automattic/jetpack-agents-manager
6 */
7
8namespace Automattic\Jetpack\Agents_Manager;
9
10use Automattic\Jetpack\Connection\Client;
11use Automattic\Jetpack\Connection\Manager as Connection_Manager;
12use Jetpack_Options;
13
14/**
15 * Class WP_REST_Jetpack_AI_JWT.
16 */
17class WP_REST_Jetpack_AI_JWT extends \WP_REST_Controller {
18
19    /**
20     * WP_REST_Jetpack_AI_JWT constructor.
21     */
22    public function __construct() {
23        $this->namespace = 'jetpack/v4';
24        $this->rest_base = '/jetpack-ai-jwt';
25    }
26
27    /**
28     * Register available routes.
29     */
30    public function register_rest_route() {
31        /*
32         * Check if the `jetpack/v4/jetpack-ai-jwt` endpoint is registered
33         * by the Jetpack plugin to avoid registering it again.
34         * In case it's not registered, register it
35         * to make it available for Jetpack products that depend on it.
36         */
37        if ( $this->is_rest_endpoint_registered() ) {
38            return;
39        }
40
41        register_rest_route(
42            $this->namespace,
43            $this->rest_base,
44            array(
45                array(
46                    'methods'             => \WP_REST_Server::EDITABLE,
47                    'callback'            => array( $this, 'get_jwt' ),
48                    'permission_callback' => array( $this, 'permission_callback' ),
49                ),
50            )
51        );
52    }
53
54    /**
55     * Check if this REST endpoint is already registered.
56     *
57     * @return bool True if the endpoint is registered, false otherwise.
58     */
59    private function is_rest_endpoint_registered() {
60        $server        = rest_get_server();
61        $routes        = $server->get_routes();
62        $full_endpoint = '/' . trim( $this->namespace, '/' ) . $this->rest_base;
63
64        return isset( $routes[ $full_endpoint ] );
65    }
66
67    /**
68     * Permission callback for the JWT endpoint.
69     *
70     * @return bool
71     */
72    public function permission_callback() {
73        return ( new Connection_Manager( 'jetpack' ) )->is_user_connected() && current_user_can( 'edit_posts' );
74    }
75
76    /**
77     * Ask WPCOM for a JWT token to use for OpenAI completion.
78     */
79    public function get_jwt() {
80        $blog_id = Jetpack_Options::get_option( 'id' );
81
82        $response = Client::wpcom_json_api_request_as_user(
83            "/sites/$blog_id/jetpack-openai-query/jwt",
84            '2',
85            array(
86                'method'  => 'POST',
87                'headers' => array( 'Content-Type' => 'application/json; charset=utf-8' ),
88            ),
89            wp_json_encode( array(), JSON_UNESCAPED_SLASHES | JSON_UNESCAPED_UNICODE ),
90            'wpcom'
91        );
92
93        if ( is_wp_error( $response ) ) {
94            return $response;
95        }
96
97        $json = json_decode( wp_remote_retrieve_body( $response ) );
98
99        if ( ! isset( $json->token ) ) {
100            return new \WP_Error(
101                'no-token',
102                'No token returned from WPCOM',
103                array( 'status' => 500 )
104            );
105        }
106
107        return rest_ensure_response(
108            array(
109                'token'   => $json->token,
110                'blog_id' => $blog_id,
111            )
112        );
113    }
114}