From fffa772c8253cff0602984b16640122a395269cf Mon Sep 17 00:00:00 2001
From: SquirrelDevelopper <squirrel@squirrelserversmanager.io>
Date: Fri, 31 Jan 2025 11:54:09 +0100
Subject: [PATCH] Enhance SFTP file download to support directories as tarballs

Updated the SFTP file handling logic to differentiate between files and directories. Added functionality to create and download tarballs for directories, ensuring efficient transfer. Improved error handling and logging for better debugging and user feedback during file operations.
---
 .../SFTPDrawer/Actions/DownloadModal.tsx      |  8 +--
 server/src/modules/ssh/SFTPInstance.ts        | 72 +++++++++++++++++--
 2 files changed, 70 insertions(+), 10 deletions(-)

diff --git a/client/src/components/DeviceComponents/SFTPDrawer/Actions/DownloadModal.tsx b/client/src/components/DeviceComponents/SFTPDrawer/Actions/DownloadModal.tsx
index 277d3c13..d9132ebe 100644
--- a/client/src/components/DeviceComponents/SFTPDrawer/Actions/DownloadModal.tsx
+++ b/client/src/components/DeviceComponents/SFTPDrawer/Actions/DownloadModal.tsx
@@ -76,15 +76,15 @@ const DownloadModal = React.forwardRef<
     fileBufferRef.current.push(chunk); // Append chunk directly to the mutable ref
   };
 
-  const handleError = (error: string) => {
-    message.error({
-      content: `Error during file download: ${error}`,
+  const handleError = (error: any) => {
+    void message.error({
+      content: `Error during file download: ${error?.message || error}`,
       duration: 6,
     });
   };
 
   const handleNotfound = (error: string) => {
-    message.error({
+    void message.error({
       content: `File not found: ${error}`,
       duration: 6,
     });
diff --git a/server/src/modules/ssh/SFTPInstance.ts b/server/src/modules/ssh/SFTPInstance.ts
index 21ae85d2..2df45e0d 100644
--- a/server/src/modules/ssh/SFTPInstance.ts
+++ b/server/src/modules/ssh/SFTPInstance.ts
@@ -1,4 +1,5 @@
 import * as util from 'node:util';
+import path from 'path';
 import { Logger } from 'pino';
 import { Socket } from 'socket.io';
 import { API, SsmEvents } from 'ssm-shared-lib';
@@ -309,14 +310,73 @@ export default class SFTPInstance {
       }
       const localRootPath = `/tmp/${v4()}`;
       FileSystemManager.createDirectory(localRootPath);
-      const localPath = `${localRootPath}/${data.path.split('/').pop()}`;
-      sftp.fastGet(data.path, localPath, (err) => {
+      const remotePath = data.path;
+      const localPath = path.join(localRootPath, path.basename(remotePath));
+
+      // Check if the remote path is a directory or a file
+      sftp.stat(remotePath, (err, stats) => {
         if (err) {
-          this.logger.error(`Error downloading file: ${err.message}`);
-          this.socket.emit(SsmEvents.SFTP.DOWNLOAD, { status: 'ERROR', message: err.message });
+          this.logger.error(`Error getting file stats: ${err.message}`);
+          this.socket.emit(SsmEvents.FileTransfer.ERROR, { status: 'ERROR', message: err.message });
         } else {
-          this.logger.info(`File downloaded from ${data.path} to ${localPath}`);
-          sendFile(this.socket, localRootPath, data.path.split('/').pop() as string);
+          if (stats.isDirectory()) {
+            const remoteTempDir = `/tmp/${v4()}`; // A unique
+            const tarPath = `${remoteTempDir}/${path.basename(remotePath)}.tar.gz`;
+            // Create a tarball from the remote directory
+            const tarCommand = `mkdir -p "${remoteTempDir}" && tar -czvf "${tarPath}" "${remotePath}"`;
+            this.logger.info(`Creating tarball: ${tarCommand}`);
+            this.sshConnectionInstance.ssh.exec(tarCommand, (err, stream) => {
+              if (err) {
+                this.logger.error(`Error creating tarball: ${err.message}`);
+                this.socket.emit(SsmEvents.FileTransfer.ERROR, {
+                  status: 'ERROR',
+                  message: err.message,
+                });
+              } else {
+                stream
+                  .on('close', (code, signal) => {
+                    if (code !== 0) {
+                      this.logger.error(`Tarball creation failed with code ${code}`);
+                      this.socket.emit(SsmEvents.FileTransfer.ERROR, {
+                        status: 'ERROR',
+                        message: `Tarball creation failed with code ${code}`,
+                      });
+                      return;
+                    }
+                    this.logger.info(`Tarball created at ${tarPath}`);
+                    sftp.fastGet(tarPath, `${localPath}.tar.gz`, (err) => {
+                      if (err) {
+                        this.logger.error(err);
+                        this.logger.error(`Error downloading file: ${err.message}`);
+                        this.socket.emit(SsmEvents.FileTransfer.ERROR, {
+                          status: 'ERROR',
+                          message: err.message,
+                        });
+                      } else {
+                        this.logger.info(`File downloaded from ${tarPath} to ${localPath}`);
+                        sendFile(this.socket, localRootPath, tarPath.split('/').pop() as string);
+                      }
+                    });
+                  })
+                  .on('data', (data) => {
+                    this.logger.info(`Tarball creation progress...`);
+                  });
+              }
+            });
+          } else {
+            sftp.fastGet(data.path, localPath, (err) => {
+              if (err) {
+                this.logger.error(`Error downloading file: ${err.message}`);
+                this.socket.emit(SsmEvents.FileTransfer.ERROR, {
+                  status: 'ERROR',
+                  message: err.message,
+                });
+              } else {
+                this.logger.info(`File downloaded from ${data.path} to ${localPath}`);
+                sendFile(this.socket, localRootPath, data.path.split('/').pop() as string);
+              }
+            });
+          }
         }
       });
     } catch (error: any) {