From 3b7cf8af571d0baeda3c1fc7dca912eb324da75f Mon Sep 17 00:00:00 2001
From: Levi Yan <leviyanx@gmail.com>
Date: Tue, 30 Jul 2024 23:35:08 +0800
Subject: [PATCH] refactor: only successfully finished first connect then open
 remote folder

---
 src/home.ts             |  9 ++++--
 src/remote-container.ts | 70 ++++++++++++++++++++++-------------------
 2 files changed, 44 insertions(+), 35 deletions(-)

diff --git a/src/home.ts b/src/home.ts
index 5bc81d8..a6c849a 100644
--- a/src/home.ts
+++ b/src/home.ts
@@ -31,8 +31,13 @@ export default class DSHome {
       async (message) => {
         switch (message.command) {
           case 'firstOpenRemoteFolder':
-            await this.remoteContainer.firstConnect(message.host, message.username, message.password, message.port);
-            this.remoteContainer.openRemoteFolder(message.host, message.path);
+            await this.remoteContainer.firstConnect(message.host, message.username, message.password, message.port)
+              .then((result) => {
+                if (result === 'success') {
+                  //  only success then open folder
+                  this.remoteContainer.openRemoteFolder(message.host, message.path);
+                }
+            })
             return;
           case 'openRemoteFolder':
             this.remoteContainer.openRemoteFolder(message.host, message.path);
diff --git a/src/remote-container.ts b/src/remote-container.ts
index 15c38f7..075a3e6 100644
--- a/src/remote-container.ts
+++ b/src/remote-container.ts
@@ -3,45 +3,49 @@ import * as path from 'path';
 import * as os from 'os';
 import * as vscode from 'vscode';
 const { NodeSSH } = require('node-ssh')
+import * as utils from './utils';
 
 export default class RemoteContainer {
 
-  async firstConnect(host: string, username: string, password: string, port: number): Promise<void> {
-    // connect the host and add the public key to the remote authorized_keys
-    const defaultPublicKeyPath = path.join(os.homedir(), '.ssh', 'id_rsa.pub');
-    // TODO: if there is no public key, generate one
-    const publicKey = fs.readFileSync(defaultPublicKeyPath, 'utf8');
+  async firstConnect(host: string, username: string, password: string, port: number): Promise<string> {
+    return new Promise(async (resolve, reject) => {
+      // connect the host and add the public key to the remote authorized_keys
+      const defaultPublicKeyPath = path.join(os.homedir(), '.ssh', 'id_rsa.pub');
+      // TODO: if there is no public key, generate one
+      const publicKey = fs.readFileSync(defaultPublicKeyPath, 'utf8');
 
+      // first connect to the remote host using password
+      const ssh = new NodeSSH();
+      try {
+        await ssh.connect({
+          host: host,
+          username: username,
+          password: password,
+          port: port
+        });
 
-    // first connect to the remote host using password
-    const ssh = new NodeSSH();
-    try {
-      await ssh.connect({
-        host: host,
-        username: username,
-        password: password,
-        port: port
-      });
+        // add the public key to the remote authorized_keys
+        await ssh.execCommand(`mkdir -p ~/.ssh && echo '${publicKey}' >> ~/.ssh/authorized_keys`);
+        console.log('Public key added to remote authorized_keys');
+        await ssh.dispose();
 
-      // add the public key to the remote authorized_keys
-      await ssh.execCommand(`mkdir -p ~/.ssh && echo '${publicKey}' >> ~/.ssh/authorized_keys`);
-      console.log('Public key added to remote authorized_keys');
-    } catch (error) {
-      console.error('Error adding public key: ', error);
-    } finally {
-      await ssh.dispose();
-    }
-
-    // TODO: request change password
-
-    // TODO: use remote.SSH.configFile
-    const sshConfigPath = path.join(os.homedir(), '.ssh', 'config');
-    // append the host to the local ssh config file 
-    const sshConfigContent =
-      `\nHost ${host}\n  HostName ${host}\n  Port ${port}\n  User ${username}\n  PreferredAuthentications publickey\n  IdentityFile ~/.ssh/id_rsa\n  `;
-    // append the host to the local ssh config file
-    fs.writeFileSync(sshConfigPath, sshConfigContent, { encoding: 'utf8', flag: 'a' });
-    console.log('Host registered in local ssh config');
+        // only connect successfully then save the host info
+        // TODO: request change password
+        // TODO: use remote.SSH.configFile
+        const sshConfigPath = path.join(os.homedir(), '.ssh', 'config');
+        // append the host to the local ssh config file 
+        const sshConfigContent =
+          `\nHost ${host}\n  HostName ${host}\n  Port ${port}\n  User ${username}\n  PreferredAuthentications publickey\n  IdentityFile ~/.ssh/id_rsa\n  `;
+        // append the host to the local ssh config file
+        fs.writeFileSync(sshConfigPath, sshConfigContent, { encoding: 'utf8', flag: 'a' });
+        console.log('Host registered in local ssh config');
+        resolve("success");
+      } catch (error) {
+        console.error('Error adding public key: ', error);
+        await ssh.dispose();
+        reject("failed");
+      }
+    });
   }