alwaysai/src/subcommands/device/connect.ts

The alwaysAI command-line interface (CLI)

import {
  CliLeaf,
  CliNumberInput,
  CliStringInput,
  CliUsageError
} from '@alwaysai/alwayscli';
import * as chalk from 'chalk';
import { Choice } from 'prompts';
import { yesCliInput } from '../../cli-inputs';
import { checkUserIsLoggedInComponent } from '../../components/user';
import {
  RequiredWithYesMessage,
  echo,
  logger,
  promptForInput
} from '../../util';
import {
  DockerStartParams,
  OpenTunnelResponseBody,
  getDockerImageName,
  getSourceAccessToken,
  processOpenTunnelData,
  removeLocalhostPortsFromKnownHostsFile,
  secureTunnelJsonFile,
  startDockerLocalProxyContainer,
  startSecureTunnelSshSession,
  stopDockerLocalProxyContainer
} from '../../util/secure-tunnel';
import { getDeviceList } from './list';
import {
  updateShadow,
  getSystemInfoShadow,
  getShadow,
  SecureTunnelPorts,
  HttpStatusCode
} from '../../util/shadows';
import { gte } from 'semver';
import * as logSymbols from 'log-symbols';

// ----------------------------------------------------------------------------
// Local types and interface
// ----------------------------------------------------------------------------
type ConnectedDevice = {
  uuid: string;
  friendlyName: string;
};
// ----------------------------------------------------------------------------
// Local functions
// ----------------------------------------------------------------------------
/**
 * Gets list of connected devices (uuid and device friendly name only)
 * @param {list} deviceList - list connected devices
 * @returns {ConnectedDevice[]} filters out disconnected devices and returns a list of connected devices
 */

const ipPortRegex =
  /^((25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?):([0-9]{1,5})$/;

function getConnectedDevices(deviceList): ConnectedDevice[] {
  const connectedDevices = deviceList.filter(
    (device) => device.status === 'connected'
  );
  return connectedDevices.map((device) => ({
    uuid: device.uuid,
    friendlyName: device.friendly_name
  }));
}

/**
 * Prompts user to select the device to connect to
 * @param {ConnectedDevice[]} connectedDevices - list of connected devices
 * @returns {string} uuid of the selected device
 */
async function selectDeviceToConnect(
  connectedDevices: ConnectedDevice[]
): Promise<string> {
  const choices: Choice[] = [{ title: 'None', value: '' }];
  connectedDevices.forEach((device) => {
    choices.push({
      title: `${device.friendlyName}\t(${device.uuid})`,
      value: `${device.uuid}`
    });
  });

  const choice = await promptForInput({
    purpose: 'to select target device to connect',
    questions: [
      {
        type: 'select',
        name: 'targetDevice',
        message: chalk.yellowBright.bold(
          'Select device to ssh to (disconnected devices are not listed)'
        ),
        initial: 0,
        choices
      }
    ]
  });

  return choice.targetDevice;
}

/**
 * Prompts user to enter user name for the remote host
 * @returns {string} user name for the remote host
 */
async function getRemoteUserLogin(): Promise<string> {
  const remoteUserLoginAnswer = await promptForInput({
    purpose: 'to set a remote user login',
    questions: [
      {
        type: 'text',
        name: 'login',
        message: chalk.yellowBright.bold('Enter a remote user login')
      }
    ]
  });
  return remoteUserLoginAnswer.login;
}

// ----------------------------------------------------------------------------
// Main function
// ----------------------------------------------------------------------------
export const deviceConnect = CliLeaf({
  name: 'connect',
  description: 'Connect to remote devices',
  namedInputs: {
    yes: yesCliInput,
    device: CliStringInput({
      description: 'Device UUID'
    }),
    'local-port': CliNumberInput({
      description: 'Local port to use for connection'
    }),
    proxy: CliStringInput({
      description:
        'HTTP proxy in the format [ip:port]. Example: 100.70.31.118:80'
    }),
    username: CliStringInput({
      description: 'Device username'
    })
  },
  hidden: true,
  async action(_, opts) {
    const { yes, device, username, proxy } = opts;
    const localPort = opts['local-port'];
    const deletedPorts = secureTunnelJsonFile().removeExpiredItems();
    await removeLocalhostPortsFromKnownHostsFile(deletedPorts);

    if (yes) {
      if (device === undefined) {
        throw new CliUsageError(RequiredWithYesMessage('device', undefined));
      }
      if (username === undefined) {
        throw new CliUsageError(RequiredWithYesMessage('username', undefined));
      }
    }

    if (proxy !== undefined && !ipPortRegex.test(proxy)) {
      throw new CliUsageError(
        'Invalid IP:Port format. Please enter a valid IP:Port'
      );
    }

    await checkUserIsLoggedInComponent({ yes });
    let selectedDeviceUuid = device;
    if (selectedDeviceUuid === undefined) {
      const deviceList = await getDeviceList();
      logger.debug(JSON.stringify(deviceList, null, 4));
      if (deviceList.length === 0) {
        return echo(chalk.redBright.bold('No devices found'));
      }

      const connectedDevices = getConnectedDevices(deviceList);
      logger.debug(
        `connectedDevices = ${JSON.stringify(connectedDevices, null, 4)}`
      );
      if (connectedDevices.length === 0) {
        return echo(chalk.redBright.bold('No connected devices found'));
      }
      selectedDeviceUuid = await selectDeviceToConnect(connectedDevices);
    }

    logger.debug(`selectedDevice = ${selectedDeviceUuid}`);
    if (selectedDeviceUuid === '') {
      return echo(chalk.redBright.bold('No device selected'));
    }

    await checkDeviceAgentSecureTunnelCompatibility(selectedDeviceUuid, 'ssh');

    //gather http ports
    let httpPorts: SecureTunnelPorts[] = [];
    if (proxy === undefined) {
      if (!yes) {
        httpPorts = await promptForHttpPortProxy();
      }
    } else {
      httpPorts.push({
        enabled: true,
        type: 'HTTP',
        ip: proxy.split(':')[0],
        port: parseInt(proxy.split(':')[1])
      });
    }

    if (httpPorts.length > 0) {
      await checkDeviceAgentSecureTunnelCompatibility(
        selectedDeviceUuid,
        'portProxy'
      );
    }

    const yesProxys = httpPorts.length > 0;
    const requiredProxys = yesProxys
      ? { SSH: 1, HTTP: 1 }
      : { SSH: 1, HTTP: 0 };
    if (yesProxys) {
      const { payload } = await constructShadow(selectedDeviceUuid, httpPorts);
      const updatedShadow = await updateShadow(selectedDeviceUuid, payload);
      logger.debug(`updatedShadow = ${JSON.stringify(updatedShadow, null, 4)}`);
    }
    const openTunnelData: OpenTunnelResponseBody = yesProxys
      ? await getSourceAccessToken(selectedDeviceUuid, requiredProxys)
      : await getSourceAccessToken(selectedDeviceUuid);
    logger.debug(`openTunnelData = ${JSON.stringify(openTunnelData, null, 4)}`);

    const secureTunnelInfoFromAWS = processOpenTunnelData(
      openTunnelData,
      selectedDeviceUuid,
      yesProxys,
      localPort
        ? (_: string) => {
            return localPort;
          }
        : undefined
    );

    if (secureTunnelInfoFromAWS.sourceAccessToken.length > 0) {
      const remoteUserLogin = username ? username : await getRemoteUserLogin();
      logger.debug(`remoteUserLogin = ${remoteUserLogin}`);

      secureTunnelJsonFile().setItem(
        selectedDeviceUuid,
        secureTunnelInfoFromAWS
      );

      const startParams: DockerStartParams = {
        deviceUuid: selectedDeviceUuid,
        imageName: getDockerImageName(),
        sshPort: secureTunnelInfoFromAWS.sshPort,
        httpPorts,
        sourceAccessToken: secureTunnelInfoFromAWS.sourceAccessToken
      };

      await startDockerLocalProxyContainer(startParams);

      await startSecureTunnelSshSession(
        remoteUserLogin,
        secureTunnelInfoFromAWS.sshPort
      );
      await removeLocalhostPortsFromKnownHostsFile([
        secureTunnelInfoFromAWS.sshPort
      ]);
      await stopDockerLocalProxyContainer(selectedDeviceUuid);
      if (proxy || yesProxys) {
        const { payload } = await constructShadow(
          selectedDeviceUuid,
          httpPorts
        );
        await disableAllPortsInShadow(
          payload.state.desired.st_ports,
          selectedDeviceUuid
        );
      }
    }
  }
});

async function promptForHttpPortProxy() {
  const answers = await promptForInput({
    purpose: 'for http proxy initialization',
    questions: [
      {
        type: 'confirm',
        name: 'value',
        message: chalk.yellowBright.bold('Add HTTP port proxy?'),
        initial: false
      },
      {
        type: (prev) => (prev ? 'text' : null),
        name: 'port',
        message: chalk.yellowBright.bold('Enter the proxy IP address:port'),
        validate: (value) => {
          return ipPortRegex.test(value)
            ? true
            : 'Invalid IP:Port format. Please enter a valid IP:Port';
        }
      }
    ]
  });
  const httpPorts: SecureTunnelPorts[] = [];
  if (answers.port) {
    httpPorts.push({
      enabled: true,
      type: 'HTTP',
      ip: answers.port.split(':')[0],
      port: parseInt(answers.port.split(':')[1])
    });
  }
  return httpPorts;
}

export async function constructShadow(
  selectedDeviceUuid: string,
  httpPorts: SecureTunnelPorts[]
) {
  const shadow = await getShadow(selectedDeviceUuid);

  let httpPortsOutput: SecureTunnelPorts[] = [];
  if (
    shadow.payload.state?.reported &&
    Object.keys(shadow.payload.state.reported).length > 0
  ) {
    const { reported } = shadow.payload.state;
    // eslint-disable-next-line prefer-const
    let { foundHttp, httpPortsInput } = updatePorts(reported, httpPorts);
    //if the exact same http proxy is not found, add it
    if (!foundHttp) {
      httpPortsInput = addHttpProxies(httpPorts, httpPortsInput);
    }
    httpPortsOutput = [...httpPortsInput];
  } else {
    httpPortsOutput = [
      ...httpPorts,
      { enabled: true, type: 'SSH', ip: '0.0.0.0', port: 22 }
    ];
  }

  const payload = {
    state: {
      desired: { st_ports: httpPortsOutput }
    },
    ...(shadow && { version: shadow.payload.version })
  };

  return { payload };
}

async function disableAllPortsInShadow(
  stPorts: SecureTunnelPorts[],
  selectedDeviceUuid: string
) {
  const shadow = await getShadow(selectedDeviceUuid);

  const disabledPorts = stPorts.map((port) => {
    port.enabled = false;
    return port;
  });
  const payload = {
    state: {
      desired: { st_ports: disabledPorts }
    },
    ...(shadow.payload.code === HttpStatusCode.OK && {
      version: shadow.payload.version
    })
  };

  await updateShadow(selectedDeviceUuid, payload);
}

export function updatePorts(reported: any, httpPorts: any[]) {
  let found = false;
  const httpPortsInput = (reported?.st_ports || []).map((port) => {
    httpPorts.forEach((ip) => {
      if (port.ip === ip.ip && port.port === ip.port) {
        port.enabled = true;
        found = true;
      } else if (port.type === 'HTTP' && port.enabled === true) {
        port.enabled = false;
      }
    });
    if (port.type === 'SSH') {
      port.enabled = true;
    }
    return port;
  });
  return { httpPortsInput, foundHttp: found };
}

export function addHttpProxies(
  httpPorts: SecureTunnelPorts[],
  httpPortsInput: { enabled: boolean; type: string; ip: string; port: number }[]
) {
  httpPorts.forEach((ip) => {
    const proxy = { enabled: true, type: 'HTTP', ip: ip.ip, port: ip.port };
    httpPortsInput.push(proxy);
  });
  return httpPortsInput;
}
export type Feature = 'ssh' | 'portProxy';

export async function checkDeviceAgentSecureTunnelCompatibility(
  thingId: string,
  feature: Feature
) {
  const shadow = await getSystemInfoShadow(thingId);

  if (shadow.versions?.agent) {
    const agentVersion = shadow.versions.agent;
    if (
      !IsFeatureSupported(SecureTunnelFeatureVersion[feature], agentVersion)
    ) {
      throw new CliUsageError(
        `The device agent version ${agentVersion} is not compatible with the feature: ${feature}. Please update the device agent to the minimum version ${SecureTunnelFeatureVersion[feature]}.`
      );
    }
  } else {
    echo(
      `${logSymbols.warning} Cannot determine ${feature} support for device ${thingId}. Ensure your CLI and Device Agent are updated to the latest versions if connection fails.`
    );
  }
}

enum SecureTunnelFeatureVersion {
  ssh = '0.2.0',
  portProxy = '1.4.0'
}

function IsFeatureSupported(
  featureVersion: string,
  currentVersion: string
): boolean {
  if (!featureVersion || !currentVersion) {
    return false;
  }
  return gte(currentVersion, featureVersion);
}