From 9a24d7bc27aec62e516ee976f047fd3db472ca0c Mon Sep 17 00:00:00 2001
From: rcooke-warwick <ryan@balena.io>
Date: Thu, 29 Jun 2023 15:36:44 +0100
Subject: [PATCH] make configurable, and maintain backwards comp

Change-type: patch
Signed-off-by: Ryan Cooke <ryan@balena.io>
---
 core/lib/common/worker.js         | 22 ++++++----------------
 core/lib/components/balena/sdk.js | 10 +++++++---
 2 files changed, 13 insertions(+), 19 deletions(-)

diff --git a/core/lib/common/worker.js b/core/lib/common/worker.js
index 3c2cf86d3..41b6cdd99 100644
--- a/core/lib/common/worker.js
+++ b/core/lib/common/worker.js
@@ -79,6 +79,7 @@ module.exports = class Worker {
 		url,
 		username,
 		sshKey,
+		sshConfig = {}
 	) {
 		this.deviceType = deviceType;
 		this.url = url;
@@ -96,26 +97,15 @@ module.exports = class Worker {
 			this.url.includes(`worker`)
 			|| this.url.includes('unix:')
 		);
-		if (this.url.includes(`balena-devices.com`)) {
-			// worker is a testbot connected to balena cloud - we ssh into it via the vpn
-			this.uuid = this.url.match(
-				/(?<=https:\/\/)(.*)(?=.balena-devices.com)/,
-			)[1];
-			this.workerHost = `ssh.balena-devices.com`;
-			this.workerUser = this.username;
-			this.workerPort = '22';
-			this.sshPrefix = `host ${this.uuid} `;
-		}
-		if (this.url.includes(`bm.balena-dev.com`)) {
-			// worker is a testbot connected to balena cloud - we ssh into it via the vpn
+		if(!this.directConnect){
 			this.uuid = this.url.match(
 				/https:\/\/([^\.]+)\./,
 			)[1];
-			this.workerHost = `ssh.devices.bm.balena-dev.com`;
-			this.workerUser = this.username;
-			this.workerPort = '222';
 			this.sshPrefix = `host ${this.uuid} `;
-		}
+			this.workerUser = this.username;
+			this.workerPort = sshConfig.port || 22;
+			this.workerHost = sshConfig.host || 'ssh.balena-devices.com'
+		}	
 	}
 
 	/**
diff --git a/core/lib/components/balena/sdk.js b/core/lib/components/balena/sdk.js
index 5f8332146..40a5b378c 100644
--- a/core/lib/components/balena/sdk.js
+++ b/core/lib/components/balena/sdk.js
@@ -65,12 +65,17 @@ module.exports = class BalenaSDK {
 	constructor(
 		apiUrl,
 		logger = { log: console.log, status: console.log, info: console.log },
+		sshConfig = {}
 	) {
 		this.balena = getSdk({
 			apiUrl: `https://api.${apiUrl}`,
 		});
 		this.pine = this.balena.pine;
 		this.logger = logger;
+		this.sshConfig = {
+			host: sshConfig.host || 'ssh.balena-devices.com',
+			port: sshConfig.port || 22
+		}
 	}
 
 	/**
@@ -91,7 +96,6 @@ module.exports = class BalenaSDK {
 			tries: 600,
 		},
 	) {
-		const sshPort = 22;
 
 		return retry(
 			async () => {
@@ -102,9 +106,9 @@ module.exports = class BalenaSDK {
 				const result = await utils.executeCommandOverSSH(
 					`host -s ${device} source /etc/profile ; ${command}`,
 					{
-						host: `ssh.${await this.balena.settings.get('proxyUrl')}`,
+						host: this.sshConfig.host,
 						username: await this.balena.auth.whoami(),
-						port: sshPort,
+						port: this.sshConfig.sshPort,
 					},
 				);