pwcli/src/commands/aws/ecs/sshtestserver/sshtestserver.ts

import { GluegunMenuToolbox } from '@lenne.tech/gluegun-menu';
import chalk = require('chalk');
import { defaultMenuSettings } from '../../../../globals';
import { getTestStacksInfo } from '../../../../common/testservers';
const { AutoComplete } = require('enquirer');

module.exports = {
  name: 'ssh-to-test-server',
  alias: ['sts'],
  description: 'SSH into a test server in ECS (a e sts)',
  hidden: false,
  run: async (toolbox: GluegunMenuToolbox) => {
    const { system, strings, print } = toolbox;

    try {
      // Get all test stacks
      print.info('Fetching available test servers...');
      const testStacks = await getTestStacksInfo(toolbox);

      if (testStacks.length === 0) {
        print.error('No test servers found');
        return;
      }

      // Filter only running/complete stacks
      const availableStacks = testStacks.filter((stack) =>
        ['CREATE_COMPLETE', 'UPDATE_COMPLETE'].includes(stack.status),
      );

      if (availableStacks.length === 0) {
        print.error('No running test servers found');
        return;
      }

      // Create options for server selection
      const options = availableStacks.map((stack) => {
        const statusColor =
          stack.status === 'CREATE_COMPLETE' ||
          stack.status === 'UPDATE_COMPLETE'
            ? chalk.green('✓')
            : chalk.yellow('⚠');
        return {
          name: `${statusColor} ${stack.name} (${stack.branch})`,
          value: stack.name,
        };
      });

      // Let user choose a server
      const prompt = new AutoComplete({
        name: 'testserver',
        message: 'Choose a test server to SSH into:',
        choices: options,
        limit: Math.min(options.length, 10),
      });

      const selectedStackName = await prompt.run();
      const selectedStack = availableStacks.find(
        (stack) => stack.name === selectedStackName,
      );

      if (!selectedStack) {
        print.error('Selected stack not found');
        return;
      }

      print.info(`Selected: ${chalk.yellow(selectedStackName)}`);

      // Get the actual ECS cluster name from CloudFormation stack
      print.info('Getting ECS cluster name from CloudFormation stack...');
      const stackDetailsResponse = await system.run(
        `aws cloudformation describe-stacks --stack-name ${selectedStackName} --output json`,
        { trim: true },
      );

      const stackDetails = JSON.parse(stackDetailsResponse);

      if (!stackDetails.Stacks || stackDetails.Stacks.length === 0) {
        print.error('Stack details not found');
        return;
      }

      const stack = stackDetails.Stacks[0];
      let clusterName = null;

      // Try to find cluster name from stack outputs first
      if (stack.Outputs) {
        const clusterOutput = stack.Outputs.find(
          (output) =>
            output.OutputKey.toLowerCase().includes('cluster') ||
            output.OutputKey.toLowerCase().includes('ecs'),
        );
        if (clusterOutput) {
          clusterName = clusterOutput.OutputValue;
        }
      }

      // If not found in outputs, try to get from stack resources
      if (!clusterName) {
        print.info(
          'Cluster name not found in outputs, checking stack resources...',
        );
        const resourcesResponse = await system.run(
          `aws cloudformation describe-stack-resources --stack-name ${selectedStackName} --output json`,
          { trim: true },
        );

        const resources = JSON.parse(resourcesResponse);
        const clusterResource = resources.StackResources.find(
          (resource) => resource.ResourceType === 'AWS::ECS::Cluster',
        );

        if (clusterResource) {
          clusterName = clusterResource.PhysicalResourceId;
        }
      }

      // If still not found, fall back to stack name (original behavior)
      if (!clusterName) {
        print.warning(
          'Could not find ECS cluster name from stack, using stack name as fallback',
        );
        clusterName = selectedStackName;
      }

      print.info(`Using ECS cluster: ${chalk.cyan(clusterName)}`);

      // Get running tasks for the cluster
      print.info('Fetching ECS tasks...');
      const tasksResponse = await system.run(
        `aws ecs list-tasks --cluster ${clusterName} --desired-status RUNNING --output json`,
        { trim: true },
      );

      const tasks = JSON.parse(tasksResponse);

      if (!tasks.taskArns || tasks.taskArns.length === 0) {
        print.error(`No running tasks found in cluster ${clusterName}`);
        return;
      }

      // Get task details to find containers
      print.info('Getting task details...');
      const taskDetailsResponse = await system.run(
        `aws ecs describe-tasks --cluster ${clusterName} --tasks ${tasks.taskArns.join(' ')} --output json`,
        { trim: true },
      );

      const taskDetails = JSON.parse(taskDetailsResponse);

      if (!taskDetails.tasks || taskDetails.tasks.length === 0) {
        print.error('No task details found');
        return;
      }

      // Find the first task with containers
      let selectedTask = null;
      let selectedContainer = null;

      for (const task of taskDetails.tasks) {
        if (task.containers && task.containers.length > 0) {
          selectedTask = task;
          // Try to find a web/app container first, otherwise use the first one
          const webContainer = task.containers.find(
            (container) =>
              container.name.toLowerCase().includes('web') ||
              container.name.toLowerCase().includes('app') ||
              container.name.toLowerCase().includes('photowall'),
          );
          selectedContainer = webContainer || task.containers[0];
          break;
        }
      }

      if (!selectedTask || !selectedContainer) {
        print.error('No suitable task/container found');
        return;
      }

      // Extract task ID from ARN
      const taskId = selectedTask.taskArn.split('/').pop();

      print.info(`Task ID: ${chalk.cyan(taskId)}`);
      print.info(`Container: ${chalk.cyan(selectedContainer.name)}`);
      print.info(`Cluster: ${chalk.cyan(clusterName)}`);

      // Build the SSH command
      const sshCommand = `aws ecs execute-command --cluster ${clusterName} --task ${taskId} --container ${selectedContainer.name} --interactive --command "/bin/bash"`;

      print.info('\nSSH command to connect to the container:');
      print.info(chalk.yellow(sshCommand));
      print.info(
        '\n' +
          chalk.green(
            'You can use this command to connect to the container manually.',
          ),
      );
      print.info(
        chalk.dim(
          'Note: Type "exit" to disconnect from the container when connected.\n',
        ),
      );

      const { confirmPaste } = await toolbox.prompt.ask([
        {
          type: 'confirm',
          name: 'confirmPaste',
          message:
            'Would you like to put this command in your clipboard for easy pasting?',
          initial: true,
        },
      ]);

      if (confirmPaste) {
        await toolbox.system.run(
          `echo "${sshCommand.replace(/"/g, '\\"')}" | pbcopy`,
        );
        print.info(chalk.green('SSH command copied to clipboard!'));
      } else {
        print.info(chalk.dim('Command not copied to clipboard.'));
      }
    } catch (error) {
      if (error.message && error.message.includes('was cancelled')) {
        print.info('Operation cancelled by user');
      } else {
        print.error(`Error: ${error.message || error}`);
        print.error('Make sure:');
        print.error('1. AWS CLI is configured and you have proper permissions');
        print.error('2. The ECS cluster exists and has running tasks');
        print.error('3. ECS Exec is enabled on the service/task');
      }
    }

    if (toolbox.fromMenu()) {
      await toolbox.menu.showMenu('aws ecs', defaultMenuSettings);
    }
  },
};