Add a setting to choose which roles are allowed to generate images. U…

…se this setting to limit functionality from loading and to limit access to the REST endpoint. Add a test around this new setting

includes/Classifai/Providers/OpenAI/DallE.php

@@ -58,7 +58,15 @@ public function can_register() {
 	public function register() {
 		$settings = $this->get_settings();
 
-		if ( current_user_can( 'upload_files' ) && isset( $settings['enable_image_gen'] ) && 1 === (int) $settings['enable_image_gen'] ) {
+		// Check if the current user has permission to generate images.
+		$roles      = $settings['roles'] ?? [];
+		$user_roles = wp_get_current_user()->roles ?? [];
+
+		if (
+			current_user_can( 'upload_files' )
+			&& ( ! empty( $roles ) && empty( array_diff( $user_roles, $roles ) ) )
+			&& ( isset( $settings['enable_image_gen'] ) && 1 === (int) $settings['enable_image_gen'] )
+		) {
 			add_action( 'admin_enqueue_scripts', [ $this, 'enqueue_admin_scripts' ] );
 			add_action( 'print_media_templates', [ $this, 'print_media_templates' ] );
 		}
@@ -195,6 +203,30 @@ function() {
 			]
 		);
 
+		// Get all roles that have the upload_files cap.
+		$roles = get_editable_roles() ?? [];
+		$roles = array_filter(
+			$roles,
+			function( $role ) {
+				return isset( $role['capabilities'], $role['capabilities']['upload_files'] ) && $role['capabilities']['upload_files'];
+			}
+		);
+		$roles = array_combine( array_keys( $roles ), array_column( $roles, 'name' ) );
+
+		add_settings_field(
+			'roles',
+			esc_html__( 'Allowed roles', 'classifai' ),
+			[ $this, 'render_checkbox_group' ],
+			$this->get_option_name(),
+			$this->get_option_name(),
+			[
+				'label_for'      => 'roles',
+				'options'        => $roles,
+				'default_values' => $default_settings['roles'],
+				'description'    => __( 'Choose which roles are allowed to generate images. Note that the roles above only include those that have permissions to upload media.', 'classifai' ),
+			]
+		);
+
 		add_settings_field(
 			'number',
 			esc_html__( 'Number of images', 'classifai' ),
@@ -247,6 +279,10 @@ public function sanitize_settings( $settings ) {
 			$new_settings['enable_image_gen'] = '1';
 		}
 
+		if ( isset( $settings['roles'] ) && is_array( $settings['roles'] ) ) {
+			$new_settings['roles'] = array_map( 'sanitize_text_field', $settings['roles'] );
+		}
+
 		if ( isset( $settings['number'] ) && is_numeric( $settings['number'] ) && (int) $settings['number'] >= 1 && (int) $settings['number'] <= 10 ) {
 			$new_settings['number'] = absint( $settings['number'] );
 		} else {
@@ -279,6 +315,7 @@ private function get_default_settings() {
 			'authenticated'    => false,
 			'api_key'          => '',
 			'enable_image_gen' => false,
+			'roles'            => array_keys( get_editable_roles() ?? [] ),
 			'number'           => 1,
 			'size'             => '1024x1024',
 		];
@@ -302,6 +339,7 @@ public function get_provider_debug_information( $settings = null, $configured =
 		return [
 			__( 'Authenticated', 'classifai' )    => $authenticated ? __( 'yes', 'classifai' ) : __( 'no', 'classifai' ),
 			__( 'Generate images', 'classifai' )  => $enabled ? __( 'yes', 'classifai' ) : __( 'no', 'classifai' ),
+			__( 'Allowed roles', 'classifai' )    => implode( ', ', $settings['roles'] ?? [] ),
 			__( 'Number of images', 'classifai' ) => absint( $settings['number'] ?? 1 ),
 			__( 'Image size', 'classifai' )       => sanitize_text_field( $settings['size'] ?? '1024x1024' ),
 			__( 'Latest response', 'classifai' )  => $this->get_formatted_latest_response( 'classifai_openai_dalle_latest_response' ),

includes/Classifai/Services/ImageProcessing.php

@@ -321,6 +321,14 @@ public function generate_image_permissions_check() {
 			return new WP_Error( 'not_enabled', esc_html__( 'Image generation not currently enabled.', 'classifai' ) );
 		}
 
+		// Check if the current user's role is allowed.
+		$roles      = $settings['roles'] ?? [];
+		$user_roles = wp_get_current_user()->roles ?? [];
+
+		if ( empty( $roles ) || ! empty( array_diff( $user_roles, $roles ) ) ) {
+			return false;
+		}
+
 		return true;
 	}
 

tests/cypress/integration/image-processing.test.js

@@ -118,6 +118,7 @@ describe('Image processing Tests', () => {
 		cy.get( '#api_key' ).clear().type( 'password' );
 
 		cy.get( '#enable_image_gen' ).check();
+		cy.get( '#openai_dalle_roles_administrator' ).check();
 		cy.get( '#number' ).select( '2' );
 		cy.get( '#size' ).select( '512x512' );
 		cy.get( '#submit' ).click();
@@ -183,7 +184,55 @@ describe('Image processing Tests', () => {
 
 		// Create test post.
 		cy.createPost( {
-			title: 'Test DALL-E post',
+			title: 'Test DALL-E post disabled',
+			content: 'Test content',
+		} );
+
+		// Close post publish panel.
+		const closePanelSelector = 'button[aria-label="Close panel"]';
+		cy.get( 'body' ).then( ( $body ) => {
+			if ( $body.find( closePanelSelector ).length > 0 ) {
+				cy.get( closePanelSelector ).click();
+			}
+		} );
+
+		// Open post settings sidebar.
+		cy.openDocumentSettingsSidebar();
+
+		// Find and open the Featured image panel.
+		const panelButtonSelector = `.components-panel__body .components-panel__body-title button:contains("Featured image")`;
+
+		cy.get( panelButtonSelector ).then( ( $panelButton ) => {
+			// Find the panel container.
+			const $panel = $panelButton.parents( '.components-panel__body' );
+
+			// Open panel.
+			if ( ! $panel.hasClass( 'is-opened' ) ) {
+				cy.wrap( $panelButton ).click();
+			}
+
+			// Click to open media modal.
+			cy.wrap( $panel )
+				.find( '.editor-post-featured-image__toggle' )
+				.click();
+
+			// Verify tab doesn't exist.
+			cy.get( '#menu-item-generate' ).should( 'not.exist' );
+		} );
+	} );
+
+	it( 'Can disable image generation by role', () => {
+		cy.visit(
+			'/wp-admin/admin.php?page=image_processing&tab=openai_dalle'
+		);
+
+		cy.get( '#enable_image_gen' ).check();
+		cy.get( '#openai_dalle_roles_administrator' ).uncheck();
+		cy.get( '#submit' ).click();
+
+		// Create test post.
+		cy.createPost( {
+			title: 'Test DALL-E post admin disabled',
 			content: 'Test content',
 		} );
 

tests/cypress/integration/language-processing.test.js

@@ -250,14 +250,14 @@ describe('Language processing Tests', () => {
 	it( 'Can disable excerpt generation feature by role', () => {
 		cy.visit( '/wp-admin/admin.php?page=language_processing&tab=openai_chatgpt' );
 
-		// Disable editor role.
+		// Disable admin role.
 		cy.get( '#enable_excerpt' ).check();
 		cy.get( '#openai_chatgpt_roles_administrator' ).uncheck();
 		cy.get( '#submit' ).click();
 
 		// Create test post.
 		cy.createPost( {
-			title: 'Test ChatGPT post editor disabled',
+			title: 'Test ChatGPT post admin disabled',
 			content: 'Test GPT content',
 		} );