migrate ssm use cases

Author: Hweinstock

package-lock.json

@@ -70,6 +70,7 @@
         "webpack-merge": "^5.10.0"
     },
     "dependencies": {
+        "@aws-sdk/client-ssm": "^3.699.0",
         "@aws-sdk/protocol-http": "^3.370.0",
         "@types/node": "^22.7.5",
         "vscode-nls": "^5.2.0",

package.json

@@ -3,13 +3,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import * as vscode from 'vscode'
-import { Session } from 'aws-sdk/clients/ssm'
-import { EC2, IAM, SSM } from 'aws-sdk'
+import { EC2, IAM } from 'aws-sdk'
 import { Ec2Selection } from './prompter'
 import { getOrInstallCli } from '../../shared/utilities/cliUtils'
 import { isCloud9 } from '../../shared/extensionUtilities'
 import { ToolkitError } from '../../shared/errors'
-import { SsmClient } from '../../shared/clients/ssmClient'
+import { SSMWrapper } from '../../shared/clients/ssm'
 import { Ec2Client } from '../../shared/clients/ec2Client'
 import {
     VscodeRemoteConnection,
@@ -29,17 +28,18 @@ import { SshConfig } from '../../shared/sshConfig'
 import { SshKeyPair } from './sshKeyPair'
 import { Ec2SessionTracker } from './remoteSessionManager'
 import { getEc2SsmEnv } from './utils'
+import { Session, StartSessionCommandOutput } from '@aws-sdk/client-ssm'
 
 export type Ec2ConnectErrorCode = 'EC2SSMStatus' | 'EC2SSMPermission' | 'EC2SSMConnect' | 'EC2SSMAgentStatus'
 
 export interface Ec2RemoteEnv extends VscodeRemoteConnection {
     selection: Ec2Selection
     keyPair: SshKeyPair
-    ssmSession: SSM.StartSessionResponse
+    ssmSession: StartSessionCommandOutput
 }
 
 export class Ec2Connecter implements vscode.Disposable {
-    protected ssmClient: SsmClient
+    protected ssm: SSMWrapper
     protected ec2Client: Ec2Client
     protected iamClient: DefaultIamClient
     protected sessionManager: Ec2SessionTracker
@@ -53,14 +53,14 @@ export class Ec2Connecter implements vscode.Disposable {
     )
 
     public constructor(readonly regionCode: string) {
-        this.ssmClient = this.createSsmSdkClient()
+        this.ssm = this.createSsmSdkClient()
         this.ec2Client = this.createEc2SdkClient()
         this.iamClient = this.createIamSdkClient()
-        this.sessionManager = new Ec2SessionTracker(regionCode, this.ssmClient)
+        this.sessionManager = new Ec2SessionTracker(regionCode, this.ssm)
     }
 
-    protected createSsmSdkClient(): SsmClient {
-        return new SsmClient(this.regionCode)
+    protected createSsmSdkClient(): SSMWrapper {
+        return new SSMWrapper(this.regionCode)
     }
 
     protected createEc2SdkClient(): Ec2Client {
@@ -71,7 +71,7 @@ export class Ec2Connecter implements vscode.Disposable {
         return new DefaultIamClient(this.regionCode)
     }
 
-    public async addActiveSession(sessionId: SSM.SessionId, instanceId: EC2.InstanceId): Promise<void> {
+    public async addActiveSession(sessionId: string, instanceId: EC2.InstanceId): Promise<void> {
         await this.sessionManager.addSession(instanceId, sessionId)
     }
 
@@ -139,7 +139,7 @@ export class Ec2Connecter implements vscode.Disposable {
     }
 
     private async checkForInstanceSsmError(selection: Ec2Selection): Promise<void> {
-        const isSsmAgentRunning = (await this.ssmClient.getInstanceAgentPingStatus(selection.instanceId)) === 'Online'
+        const isSsmAgentRunning = (await this.ssm.getInstanceAgentPingStatus(selection.instanceId)) === 'Online'
 
         if (!isSsmAgentRunning) {
             this.throwConnectionError('Is SSM Agent running on the target instance?', selection, {
@@ -173,15 +173,15 @@ export class Ec2Connecter implements vscode.Disposable {
             shellArgs: shellArgs,
         }
 
-        await openRemoteTerminal(terminalOptions, () => this.ssmClient.terminateSession(session)).catch((err) => {
+        await openRemoteTerminal(terminalOptions, () => this.ssm.terminateSession(session)).catch((err) => {
             throw ToolkitError.chain(err, 'Failed to open ec2 instance.')
         })
     }
 
     public async attemptToOpenEc2Terminal(selection: Ec2Selection): Promise<void> {
         await this.checkForStartSessionError(selection)
         try {
-            const response = await this.ssmClient.startSession(selection.instanceId)
+            const response = await this.ssm.startSession(selection.instanceId)
             await this.openSessionInTerminal(response, selection)
         } catch (err: unknown) {
             this.throwGeneralConnectionError(selection, err as Error)
@@ -222,7 +222,7 @@ export class Ec2Connecter implements vscode.Disposable {
 
             throw err
         }
-        const ssmSession = await this.ssmClient.startSession(selection.instanceId, 'AWS-StartSSHSession')
+        const ssmSession = await this.ssm.startSession(selection.instanceId, 'AWS-StartSSHSession')
         await this.addActiveSession(selection.instanceId, ssmSession.SessionId!)
 
         const vars = getEc2SsmEnv(selection, ssm, ssmSession)
@@ -270,13 +270,13 @@ export class Ec2Connecter implements vscode.Disposable {
         const command = `echo "${sshPubKey}" > ${remoteAuthorizedKeysPaths}`
         const documentName = 'AWS-RunShellScript'
 
-        await this.ssmClient.sendCommandAndWait(selection.instanceId, documentName, {
+        await this.ssm.sendCommandAndWait(selection.instanceId, documentName, {
             commands: [command],
         })
     }
 
     public async getRemoteUser(instanceId: string) {
-        const osName = await this.ssmClient.getTargetPlatformName(instanceId)
+        const osName = await this.ssm.getTargetPlatformName(instanceId)
         if (osName === 'Amazon Linux') {
             return 'ec2-user'
         }

packages/core/src/awsService/ec2/model.ts

@@ -3,30 +3,30 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { EC2, SSM } from 'aws-sdk'
-import { SsmClient } from '../../shared/clients/ssmClient'
+import { EC2 } from 'aws-sdk'
+import { SSMWrapper } from '../../shared/clients/ssm'
 import { Disposable } from 'vscode'
 
-export class Ec2SessionTracker extends Map<EC2.InstanceId, SSM.SessionId> implements Disposable {
+export class Ec2SessionTracker extends Map<EC2.InstanceId, string> implements Disposable {
     public constructor(
         readonly regionCode: string,
-        protected ssmClient: SsmClient
+        protected ssm: SSMWrapper
     ) {
         super()
     }
 
-    public async addSession(instanceId: EC2.InstanceId, sessionId: SSM.SessionId): Promise<void> {
+    public async addSession(instanceId: EC2.InstanceId, sessionId: string): Promise<void> {
         if (this.isConnectedTo(instanceId)) {
             const existingSessionId = this.get(instanceId)!
-            await this.ssmClient.terminateSessionFromId(existingSessionId)
+            await this.ssm.terminateSessionFromId(existingSessionId)
             this.set(instanceId, sessionId)
         } else {
             this.set(instanceId, sessionId)
         }
     }
 
     private async disconnectEnv(instanceId: EC2.InstanceId): Promise<void> {
-        await this.ssmClient.terminateSessionFromId(this.get(instanceId)!)
+        await this.ssm.terminateSessionFromId(this.get(instanceId)!)
         this.delete(instanceId)
     }
 

packages/core/src/awsService/ec2/remoteSessionManager.ts

@@ -6,7 +6,7 @@
 import assert from 'assert'
 import * as sinon from 'sinon'
 import { Ec2Connecter } from '../../../awsService/ec2/model'
-import { SsmClient } from '../../../shared/clients/ssmClient'
+import { SSMWrapper } from '../../../shared/clients/ssm'
 import { Ec2Client } from '../../../shared/clients/ec2Client'
 import { Ec2Selection } from '../../../awsService/ec2/prompter'
 import { ToolkitError } from '../../../shared/errors'
@@ -110,7 +110,7 @@ describe('Ec2ConnectClient', function () {
             sinon.stub(Ec2Connecter.prototype, 'isInstanceRunning').resolves(true)
             sinon.stub(Ec2Connecter.prototype, 'getAttachedIamRole').resolves({ Arn: 'testRole' } as IAM.Role)
             sinon.stub(Ec2Connecter.prototype, 'hasProperPermissions').resolves(true)
-            sinon.stub(SsmClient.prototype, 'getInstanceAgentPingStatus').resolves('offline')
+            sinon.stub(SSMWrapper.prototype, 'getInstanceAgentPingStatus').resolves('offline')
 
             try {
                 await client.checkForStartSessionError(instanceSelection)
@@ -124,15 +124,15 @@ describe('Ec2ConnectClient', function () {
             sinon.stub(Ec2Connecter.prototype, 'isInstanceRunning').resolves(true)
             sinon.stub(Ec2Connecter.prototype, 'getAttachedIamRole').resolves({ Arn: 'testRole' } as IAM.Role)
             sinon.stub(Ec2Connecter.prototype, 'hasProperPermissions').resolves(true)
-            sinon.stub(SsmClient.prototype, 'getInstanceAgentPingStatus').resolves('Online')
+            sinon.stub(SSMWrapper.prototype, 'getInstanceAgentPingStatus').resolves('Online')
 
             assert.doesNotThrow(async () => await client.checkForStartSessionError(instanceSelection))
         })
     })
 
     describe('sendSshKeysToInstance', async function () {
         it('calls the sdk with the proper parameters', async function () {
-            const sendCommandStub = sinon.stub(SsmClient.prototype, 'sendCommandAndWait')
+            const sendCommandStub = sinon.stub(SSMWrapper.prototype, 'sendCommandAndWait')
 
             const testSelection = {
                 instanceId: 'test-id',
@@ -146,7 +146,7 @@ describe('Ec2ConnectClient', function () {
         })
 
         it('avoids writing the keys to any telemetry metrics', async function () {
-            sinon.stub(SsmClient.prototype, 'sendCommandAndWait')
+            sinon.stub(SSMWrapper.prototype, 'sendCommandAndWait')
 
             const testSelection = {
                 instanceId: 'test-id',
@@ -168,7 +168,7 @@ describe('Ec2ConnectClient', function () {
         let getTargetPlatformNameStub: sinon.SinonStub<[target: string], Promise<string>>
 
         before(async function () {
-            getTargetPlatformNameStub = sinon.stub(SsmClient.prototype, 'getTargetPlatformName')
+            getTargetPlatformNameStub = sinon.stub(SSMWrapper.prototype, 'getTargetPlatformName')
         })
 
         after(async function () {

packages/core/src/shared/clients/ssmClient.ts

@@ -6,11 +6,11 @@
 import assert from 'assert'
 import * as sinon from 'sinon'
 import { Ec2SessionTracker } from '../../../awsService/ec2/remoteSessionManager'
-import { SsmClient } from '../../../shared/clients/ssmClient'
+import { SSMWrapper } from '../../../shared/clients/ssm'
 
 describe('Ec2RemoteSessionManager', async function () {
     it('maintains connections to instances', async function () {
-        const envManager = new Ec2SessionTracker('test-region', new SsmClient('test-region'))
+        const envManager = new Ec2SessionTracker('test-region', new SSMWrapper('test-region'))
         await envManager.addSession('test-instance', 'test-env')
         await envManager.addSession('test-instance2', 'test-env2')
         await envManager.addSession('test-instance3', 'test-env3')
@@ -22,8 +22,8 @@ describe('Ec2RemoteSessionManager', async function () {
     })
 
     it('only allows for single connection with any given instance', async function () {
-        const envManager = new Ec2SessionTracker('test-region', new SsmClient('test-region'))
-        const terminateStub = sinon.stub(SsmClient.prototype, 'terminateSessionFromId')
+        const envManager = new Ec2SessionTracker('test-region', new SSMWrapper('test-region'))
+        const terminateStub = sinon.stub(SSMWrapper.prototype, 'terminateSessionFromId')
 
         await envManager.addSession('test-instance', 'test-env')
         sinon.assert.notCalled(terminateStub)
@@ -37,8 +37,8 @@ describe('Ec2RemoteSessionManager', async function () {
     })
 
     it('closes all active connections', async function () {
-        const envManager = new Ec2SessionTracker('test-region', new SsmClient('test-region'))
-        const terminateStub = sinon.stub(SsmClient.prototype, 'terminateSessionFromId')
+        const envManager = new Ec2SessionTracker('test-region', new SSMWrapper('test-region'))
+        const terminateStub = sinon.stub(SSMWrapper.prototype, 'terminateSessionFromId')
 
         await envManager.addSession('test-instance', 'test-env')
         await envManager.addSession('test-instance2', 'test-env2')