From bc31c70c9ccde8b093f3e71a7678bdcdbb4fbf8d Mon Sep 17 00:00:00 2001
From: Sergei Parshev <parshev@adobe.com>
Date: Fri, 20 Jun 2025 13:01:23 -0400
Subject: [PATCH] Adds SSHConfigOtp for connect_ssh step for One-Time SSH
 access

---
 communicator/step_connect_ssh.go              | 16 ++++++++++----
 .../communicator/ssh/communicator.go          | 22 +++++++++++++++++--
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/communicator/step_connect_ssh.go b/communicator/step_connect_ssh.go
index 8f7939c95..3a9566ec3 100644
--- a/communicator/step_connect_ssh.go
+++ b/communicator/step_connect_ssh.go
@@ -31,10 +31,11 @@ import (
 // In general, you should use StepConnect.
 type StepConnectSSH struct {
 	// All the fields below are documented on StepConnect
-	Config    *Config
-	Host      func(multistep.StateBag) (string, error)
-	SSHConfig func(multistep.StateBag) (*gossh.ClientConfig, error)
-	SSHPort   func(multistep.StateBag) (int, error)
+	Config       *Config
+	Host         func(multistep.StateBag) (string, error)
+	SSHConfig    func(multistep.StateBag) (*gossh.ClientConfig, error)
+	SSHConfigOtp func(multistep.StateBag) (*gossh.ClientConfig, error)
+	SSHPort      func(multistep.StateBag) (int, error)
 }
 
 func (s *StepConnectSSH) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
@@ -217,6 +218,13 @@ func (s *StepConnectSSH) waitForSSH(state multistep.StateBag, ctx context.Contex
 			Timeout:                s.Config.SSHReadWriteTimeout,
 			Tunnels:                tunnels,
 		}
+		if s.SSHConfigOtp != nil {
+			// Create dynamic SSH config function
+			sshConfigFunc := func() (*gossh.ClientConfig, error) {
+				return s.SSHConfigOtp(state)
+			}
+			config.SSHConfigOtp = sshConfigFunc
+		}
 
 		log.Printf("[INFO] Attempting SSH connection to %s...", address)
 		comm, err = ssh.New(address, config)
diff --git a/sdk-internals/communicator/ssh/communicator.go b/sdk-internals/communicator/ssh/communicator.go
index 538f08a56..a65fe6c26 100644
--- a/sdk-internals/communicator/ssh/communicator.go
+++ b/sdk-internals/communicator/ssh/communicator.go
@@ -60,9 +60,13 @@ type TunnelSpec struct {
 
 // Config is the structure used to configure the SSH communicator.
 type Config struct {
-	// The configuration of the Go SSH connection
+	// The configuration of the Go SSH connection (static fallback)
 	SSHConfig *ssh.ClientConfig
 
+	// SSHConfigOtp returns a new SSH config. This is called for each new connection
+	// to allow for dynamic credential generation (like one-time passwords)
+	SSHConfigOtp func() (*ssh.ClientConfig, error)
+
 	// Connection returns a new connection. The current connection
 	// in use will be closed as part of the Close method, or in the
 	// case an error occurs.
@@ -327,6 +331,20 @@ func (c *comm) reconnect() (err error) {
 
 	log.Printf("[DEBUG] handshaking with SSH")
 
+	// Get fresh SSH config if available, otherwise use static config
+	var sshConfig *ssh.ClientConfig
+	if c.config.SSHConfigOtp != nil {
+		log.Printf("[DEBUG] Getting fresh SSH config for connection")
+		sshConfig, err = c.config.SSHConfigOtp()
+		if err != nil {
+			log.Printf("[ERROR] Failed to get fresh SSH config: %s", err)
+			return err
+		}
+	} else {
+		log.Printf("[DEBUG] Using static SSH config")
+		sshConfig = c.config.SSHConfig
+	}
+
 	// Default timeout to 1 minute if it wasn't specified (zero value). For
 	// when you need to handshake from low orbit.
 	var duration time.Duration
@@ -343,7 +361,7 @@ func (c *comm) reconnect() (err error) {
 	var req <-chan *ssh.Request
 
 	go func() {
-		sshConn, sshChan, req, err = ssh.NewClientConn(c.conn, c.address, c.config.SSHConfig)
+		sshConn, sshChan, req, err = ssh.NewClientConn(c.conn, c.address, sshConfig)
 		close(connectionEstablished)
 	}()