Added recursive download and tab icon when a filesystem is shared (#140)

* Added recursive download

* Updated player's download documentation

* Update changelog

CHANGELOG.adoc

@@ -6,6 +6,22 @@
 This document provides a high-level view of the changes introduced in {project-name} by release.
 For a detailed view of what has changed, refer to the {uri-repo}/commits/master[commit history] on GitHub.
 
+== 0.4.0 - <unreleased>
+
+=== Enhancements
+
+* Add recursive folder download from the PyRDP Player and a queue to download files ({uri-issue}140[#140])
+
+
+=== Bug fixes
+
+=== Credits
+
+Thanks to the following people who contributed to this release:
+
+Maxime Carbonneau
+
+
 == 0.3.0 - 2019-08-31
 
 A special BlackHat USA Arsenal 2019 release!

pyrdp/player/FileDownloadDialog.py

@@ -5,30 +5,41 @@
 #
 from PySide2.QtCore import QObject, Qt
 from PySide2.QtGui import QCloseEvent
-from PySide2.QtWidgets import QDialog, QLabel, QMessageBox, QProgressBar, QVBoxLayout
+from PySide2.QtWidgets import QDialog, QLabel, QMessageBox, QProgressBar, QPushButton, QVBoxLayout
 
 
 class FileDownloadDialog(QDialog):
-    def __init__(self, remotePath: str, targetPath: str, parent: QObject):
+    def __init__(self, remotePath: str, targetPath: str, isMultipleDownload: bool, parent: QObject):
         super().__init__(parent, Qt.CustomizeWindowHint | Qt.WindowTitleHint)
         self.titleLabel = QLabel(f"Downloading {remotePath} to {targetPath}")
 
         self.progressBar = QProgressBar()
         self.progressBar.setMinimum(0)
         self.progressBar.setMaximum(0)
+
+
         self.actualProgress = 0
         self.actualMaximum = 0
         self.isComplete = False
 
-        self.progressLabel = QLabel("0 bytes")
+        self.isMultipleDownload = isMultipleDownload
+        self.downloadCount = 0
+        self.downloadTotal = 0
+
+        self.progressLabel = QLabel(f"{self.downloadCount} / {self.downloadTotal} files downloaded")
+        self.progressSizeLabel = QLabel("0 bytes")
 
         self.widgetLayout = QVBoxLayout()
         self.widgetLayout.addWidget(self.titleLabel)
-        self.widgetLayout.addWidget(self.progressBar)
         self.widgetLayout.addWidget(self.progressLabel)
+        self.widgetLayout.addWidget(self.progressBar)
+        self.widgetLayout.addWidget(self.progressSizeLabel)
+
+        self.closeButton = QPushButton("Continue download in background")
+        self.closeButton.clicked.connect(self.hide)
+        self.widgetLayout.addWidget(self.closeButton)
 
         self.setLayout(self.widgetLayout)
-        self.setModal(True)
 
     def getHumanReadableSize(self, size: int):
         prefixes = ["", "K", "M", "G"]
@@ -42,6 +53,14 @@ def getHumanReadableSize(self, size: int):
 
         return f"{'%.2f' % size if size % 1 != 0 else int(size)} {prefixes[0]}"
 
+    def incrementDownloadCount(self):
+        self.downloadCount += 1
+        self.progressLabel.setText(f"{self.downloadCount} / {self.downloadTotal} files downloaded")
+
+    def incrementDownloadTotal(self):
+        self.downloadTotal += 1
+        self.progressLabel.setText(f"{self.downloadCount} / {self.downloadTotal} files downloaded")
+
     def updateProgress(self):
         progress = self.getHumanReadableSize(self.actualProgress)
 
@@ -50,10 +69,10 @@ def updateProgress(self):
             maximum = self.getHumanReadableSize(self.actualMaximum)
 
             self.progressBar.setValue(percentage)
-            self.progressLabel.setText(f"{progress}B / {maximum}B ({percentage}%)")
+            self.progressSizeLabel.setText(f"{progress}B / {maximum}B ({percentage}%)")
         else:
             self.progressBar.setValue(0)
-            self.progressLabel.setText(f"{progress}B")
+            self.progressSizeLabel.setText(f"{progress}B")
 
     def reportSize(self, maximum: int):
         self.actualMaximum = maximum
@@ -66,15 +85,26 @@ def reportSize(self, maximum: int):
         self.updateProgress()
 
     def reportProgress(self, progress: int):
-        self.actualProgress = progress
+        self.actualProgress += progress
         self.updateProgress()
 
     def reportCompletion(self, error: int):
+        self.show()
+
+        # QMessageBox is a modal, meaning that execution of code stops until the window closes
+        # We need our own message box if we want to queue multiple downloads
+        box = QMessageBox(self)
+        box.setWindowTitle("Download Complete")
+        box.setWindowModality(Qt.NonModal)
+
         if error == 0:
-            QMessageBox.information(self, "Download Complete", "Download completed successfully.")
+            box.setIcon(QMessageBox.Information)
+            box.setText("Download completed successfully.")
         else:
-            QMessageBox.critical(self, "Download Complete", f"Download failed. Error code: {'0x%08lx' % error}")
+            box.setIcon(QMessageBox.Critical)
+            box.setText(f"Download failed. Error code: {'0x%08lx' % error}")
 
+        box.show()
         self.isComplete = True
         self.close()
 

pyrdp/player/FileSystemWidget.py

@@ -8,20 +8,21 @@
 from typing import Optional
 
 from PySide2.QtCore import QObject, QPoint, Qt, Signal
-from PySide2.QtWidgets import QAction, QFileDialog, QFrame, QLabel, QListWidget, QMenu, QVBoxLayout, QWidget
+from PySide2.QtWidgets import QAction, QFileDialog, QFrame, QLabel, QListWidget, QMenu, QMessageBox, QVBoxLayout, QWidget
 
 from pyrdp.player.FileDownloadDialog import FileDownloadDialog
 from pyrdp.player.filesystem import Directory, DirectoryObserver, File, FileSystemItemType
 from pyrdp.player.FileSystemItem import FileSystemItem
 
+import os
 
 class FileSystemWidget(QWidget, DirectoryObserver):
     """
     Widget for listing directory contents and download files from the RDP client.
     """
 
-    # fileDownloadRequested(file, targetPath, dialog)
     fileDownloadRequested = Signal(File, str, FileDownloadDialog)
+    directoryDownloadRequested = Signal(Directory, str, FileDownloadDialog)
 
     def __init__(self, root: Directory, parent: QObject = None):
         """
@@ -99,6 +100,7 @@ def onItemDoubleClicked(self, item: FileSystemItem):
 
         self.listCurrentDirectory()
 
+
     def listCurrentDirectory(self):
         """
         Refresh the list widget with the current directory's contents.
@@ -175,8 +177,13 @@ def onCustomContextMenu(self, localPosition: QPoint):
         downloadAction.setEnabled(selectedFile.type in [FileSystemItemType.File])
         downloadAction.triggered.connect(self.downloadFile)
 
+        downloadRecursiveAction = QAction("Download files recursively")
+        downloadRecursiveAction.setEnabled(selectedFile.type in [FileSystemItemType.Directory])
+        downloadRecursiveAction.triggered.connect(self.downloadDirectoryRecursively)
+
         itemMenu = QMenu()
         itemMenu.addAction(downloadAction)
+        itemMenu.addAction(downloadRecursiveAction)
 
         itemMenu.exec_(globalPosition)
 
@@ -189,8 +196,39 @@ def downloadFile(self):
         filePath = file.getFullPath()
         targetPath, _ = QFileDialog.getSaveFileName(self, f"Download file {filePath}", file.name)
 
-        if targetPath != "":
-            dialog = FileDownloadDialog(filePath, targetPath, self)
-            dialog.show()
+        if targetPath == "":
+            QMessageBox.critical(self, "Download file", "Please select a valid file. Aborting download.")
+            return
+
+        dialog = FileDownloadDialog(filePath, targetPath, False, self)
+        dialog.incrementDownloadTotal()
+        dialog.show()
+
+        self.fileDownloadRequested.emit(file, targetPath, dialog)
+
+    def downloadDirectoryRecursively(self):
+        selectedFolder = self.selectedFile()
+
+        if selectedFolder.type != FileSystemItemType.Directory:
+            return
+
+        directoryPath = QFileDialog.getExistingDirectory(self, f"Download folder {selectedFolder.getFullPath()}")
+        remotePath = selectedFolder.getFullPath()
+
+        dialog = None
+        if directoryPath == "":
+            QMessageBox.critical(self, "Download folder", f"Please select a valid folder. Aborting download.")
+            return
+
+        directoryPath += f"/{selectedFolder.name}"
+
+        try:
+            os.mkdir(directoryPath)
+        except FileExistsError:
+            QMessageBox.critical(self, "Download folder", f"Folder already exist. Make sure to select an empty directory. Aborting download.")
+            return
+
+        dialog = FileDownloadDialog(remotePath, directoryPath, True, self)
+        dialog.show()
 
-            self.fileDownloadRequested.emit(file, targetPath, dialog)
+        self.directoryDownloadRequested.emit(selectedFolder, directoryPath, dialog)