Simplification

Author: swissspidy

src/AI_Command.php

@@ -29,7 +29,7 @@
  *     WordPress shines bright
  *
  *     # Generate an image from a prompt
- *     $ wp ai generate image "A futuristic WordPress logo" --output=logo.png
+ *     $ wp ai generate image "A futuristic WordPress logo" --destination-file=logo.png
  *     Success: Image saved to logo.png
  *
  *     # Check if a prompt is supported
@@ -40,46 +40,6 @@
  */
 class AI_Command extends WP_CLI_Command {
 
-	/**
-	 * Maximum binary image size in bytes (50MB).
-	 */
-	const MAX_IMAGE_SIZE_BYTES = 52428800; // 50 * 1024 * 1024
-
-	/**
-	 * Dummy prompt used for capability checking.
-	 *
-	 * This constant provides a consistent prompt value when checking AI capabilities.
-	 * The specific content doesn't affect capability detection, which is based on
-	 * configured providers and their available features.
-	 */
-	const CAPABILITY_CHECK_PROMPT = 'capability-check';
-
-	/**
-	 * System directories that should be protected from file writes.
-	 */
-	const FORBIDDEN_PATHS = array(
-		// Unix/Linux system directories
-		'/etc',
-		'/bin',
-		'/usr/bin',
-		'/sbin',
-		'/usr/sbin',
-		'/boot',
-		'/sys',
-		'/proc',
-		// Windows system directories (case-insensitive)
-		'C:\\Windows',
-		'C:\\Program Files',
-		'C:\\Program Files (x86)',
-	);
-
-	/**
-	 * Maximum size for base64-encoded image data.
-	 * Base64 encoding increases size by ~33%, so 50MB binary = ~67MB base64.
-	 * Using 70MB as safe upper bound.
-	 */
-	const MAX_IMAGE_SIZE_BASE64 = 70000000;
-
 	/**
 	 * Generates AI content.
 	 *
@@ -117,9 +77,12 @@ class AI_Command extends WP_CLI_Command {
 	 * [--system-instruction=<instruction>]
 	 * : System instruction to guide the AI's behavior.
 	 *
-	 * [--output=<file>]
+	 * [--destination-file=<file>]
 	 * : For image generation, path to save the generated image.
 	 *
+	 * [--stdout]
+	 * Output the whole image using standard output (incompatible with --destination-file=)
+	 *
 	 * [--format=<format>]
 	 * : Output format for text.
 	 * ---
@@ -150,7 +113,7 @@ class AI_Command extends WP_CLI_Command {
 	 *     $ wp ai generate image "A minimalist WordPress logo" --output=wp-logo.png
 	 *
 	 * @param array{0: string, 1: string} $args Positional arguments.
-	 * @param array{model: string, provider: string, temperature: float, 'top-p': float, 'top-k': int, 'max-tokens': int, 'system-instruction': string, output: string, format: string} $assoc_args Associative arguments.
+	 * @param array{model: string, provider: string, temperature: float, 'top-p': float, 'top-k': int, 'max-tokens': int, 'system-instruction': string, 'destination-file': string, stdout: bool, format: string} $assoc_args Associative arguments.
 	 * @return void
 	 */
 	public function generate( $args, $assoc_args ) {
@@ -322,7 +285,7 @@ public function status( $args, $assoc_args ) {
 
 		try {
 			// Create a builder to check capabilities (using constant for consistency)
-			$builder = AI_Client::prompt( self::CAPABILITY_CHECK_PROMPT );
+			$builder = AI_Client::prompt();
 
 			// Check each capability
 			$capabilities = array(
@@ -388,8 +351,8 @@ private function generate_text( $builder, $assoc_args ) {
 	/**
 	 * Generates an image from the prompt builder.
 	 *
-	 * @param \WordPress\AI_Client\Builders\Prompt_Builder $builder    The prompt builder.
-	 * @param array{output: string}                        $assoc_args Associative arguments.
+	 * @param \WordPress\AI_Client\Builders\Prompt_Builder         $builder    The prompt builder.
+	 * @param array{'destination-file': string, stdout: bool} $assoc_args Associative arguments.
 	 * @return void
 	 */
 	private function generate_image( $builder, $assoc_args ) {
@@ -398,75 +361,66 @@ private function generate_image( $builder, $assoc_args ) {
 			WP_CLI::error( 'Image generation is not supported. Make sure AI provider credentials are configured.' );
 		}
 
-		$image_file = $builder->generate_image();
-
-		if ( isset( $assoc_args['output'] ) ) {
-			$output_path = $assoc_args['output'];
+		if ( ! empty( $assoc_args['stdout'] ) && ! empty( $assoc_args['dir'] ) ) {
+			WP_CLI::error( '--stdout and --destination-file cannot be used together.' );
+		}
 
-			// Resolve the full real path
-			$parent_dir = dirname( $output_path );
+		if ( isset( $assoc_args['destination-file'] ) ) {
+			$output_path = $assoc_args['destination-file'];
+			$parent_dir  = dirname( $output_path );
 
-			// Check if parent directory exists
-			if ( ! file_exists( $parent_dir ) || ! is_dir( $parent_dir ) ) {
+			if ( ! is_dir( $parent_dir ) ) {
 				WP_CLI::error( 'Invalid output directory. Directory does not exist: ' . $parent_dir );
 			}
+		}
 
-			// Resolve the real path to prevent traversal attacks
-			$real_parent_dir = realpath( $parent_dir );
-			if ( false === $real_parent_dir ) {
-				WP_CLI::error( 'Cannot resolve output directory path.' );
-			}
-
-			// Reconstruct the output path with the resolved parent directory
-			$safe_output_path = $real_parent_dir . DIRECTORY_SEPARATOR . basename( $output_path );
-
-			// Prevent writing to sensitive system directories
-			foreach ( self::FORBIDDEN_PATHS as $forbidden ) {
-				// Use case-sensitive check for Unix paths (start with /), case-insensitive for Windows (contain :\)
-				$is_windows_path = ( false !== strpos( $forbidden, ':\\' ) );
-				$matches         = $is_windows_path
-					? ( 0 === stripos( $real_parent_dir, $forbidden ) )
-					: ( 0 === strpos( $real_parent_dir, $forbidden ) );
+		$image_file = $builder->generate_image();
 
-				if ( $matches ) {
-					WP_CLI::error( 'Cannot write to system directory: ' . $safe_output_path );
-				}
-			}
+		if ( isset( $assoc_args['destination-file'] ) ) {
+			$output_path = $assoc_args['destination-file'];
+			$output_path = realpath( dirname( $output_path ) ) . DIRECTORY_SEPARATOR . basename( $output_path );
 
-			// Get the image content from data URI
 			$data_uri = $image_file->getDataUri();
 
-			// Extract base64 data from data URI
 			$data_parts = $data_uri ? explode( ',', $data_uri, 2 ) : [];
 			if ( count( $data_parts ) !== 2 ) {
 				WP_CLI::error( 'Invalid image data received.' );
 			}
 
-			// Validate and decode base64 data
 			$base64_data = $data_parts[1];
 
-			// Check reasonable size limit
-			if ( strlen( $base64_data ) > self::MAX_IMAGE_SIZE_BASE64 ) {
-				WP_CLI::error( 'Image data exceeds maximum size limit (50MB).' );
-			}
-
-			// Try strict base64 decode - this validates format
 			// phpcs:ignore WordPress.PHP.DiscouragedPHPFunctions.obfuscation_base64_decode
 			$image_data = base64_decode( $base64_data, true );
 			if ( false === $image_data ) {
 				WP_CLI::error( 'Invalid base64 image data format.' );
 			}
 
 			// Save to file
-			$result = file_put_contents( $safe_output_path, $image_data );
+			$result = file_put_contents( $output_path, $image_data );
 			if ( false === $result ) {
-				WP_CLI::error( 'Failed to save image to ' . $safe_output_path );
+				WP_CLI::error( 'Failed to save image to ' . $output_path );
+			}
+
+			WP_CLI::success( 'Image saved to ' . $output_path );
+		} elseif ( $assoc_args['stdout'] ) {
+			$data_uri = $image_file->getDataUri();
+
+			$data_parts = $data_uri ? explode( ',', $data_uri, 2 ) : [];
+			if ( count( $data_parts ) !== 2 ) {
+				WP_CLI::error( 'Invalid image data received.' );
+			}
+
+			$base64_data = $data_parts[1];
+
+			// phpcs:ignore WordPress.PHP.DiscouragedPHPFunctions.obfuscation_base64_decode
+			$image_data = base64_decode( $base64_data, true );
+			if ( false === $image_data ) {
+				WP_CLI::error( 'Invalid base64 image data format.' );
 			}
 
-			WP_CLI::success( 'Image saved to ' . $safe_output_path );
+			WP_CLI::log( $image_data );
 		} else {
-			// Output data URI
-			WP_CLI::success( 'Image generated (data URI):' );
+			WP_CLI::success( 'Image generated:' );
 			WP_CLI::line( (string) $image_file->getDataUri() );
 		}
 	}