Merge pull request #119 from GoSecure/player_tabname

Pyrdp-player improvements

pyrdp/player/BaseWindow.py

@@ -5,8 +5,10 @@
 #
 
 import logging
+from typing import Dict
 
-from PySide2.QtWidgets import QTabWidget, QWidget
+from PySide2.QtGui import QKeySequence
+from PySide2.QtWidgets import QShortcut, QTabWidget, QWidget
 
 from pyrdp.logging import LOGGER_NAMES
 
@@ -17,18 +19,32 @@ class BaseWindow(QTabWidget):
     regardless of their origin (network or file).
     """
 
-    def __init__(self, parent: QWidget = None, maxTabCount = 250):
+    def __init__(self, options: Dict[str, object], parent: QWidget = None, maxTabCount = 250):
         super().__init__(parent)
         self.maxTabCount = maxTabCount
         self.setTabsClosable(True)
-        self.tabCloseRequested.connect(self.onTabClosed)
+        self.tabCloseRequested.connect(self.onTabCloseRequest)
         self.log = logging.getLogger(LOGGER_NAMES.PLAYER)
+        self.options = options
+        self.closeTabShortcut = QShortcut(QKeySequence("Ctrl+W"), self, self.onCtrlW)
 
-    def onTabClosed(self, index):
+    def onTabClosed(self, index: int):
         """
         Gracefully closes the tab by calling the onClose method
         :param index: Index of the closed tab
         """
         tab = self.widget(index)
         tab.onClose()
         self.removeTab(index)
+
+
+    def onTabCloseRequest(self, index: int):
+        """
+        By default, will close the tab. Can be overriden to add validation.
+        """
+
+        self.onTabClosed(index)
+
+    def onCtrlW(self):
+        if self.options.get("closeTabOnCtrlW") and self.count() > 0:
+            self.onTabCloseRequest(self.currentIndex())

pyrdp/player/LiveEventHandler.py

@@ -7,12 +7,15 @@
 from pathlib import Path, PosixPath
 from typing import BinaryIO, Dict
 
+from PySide2.QtCore import Signal
 from PySide2.QtWidgets import QTextEdit
 
 from pyrdp.enum import DeviceType, PlayerPDUType
 from pyrdp.layer import PlayerLayer
 from pyrdp.pdu import PlayerDeviceMappingPDU, PlayerDirectoryListingRequestPDU, PlayerDirectoryListingResponsePDU, \
-    PlayerFileDownloadCompletePDU, PlayerFileDownloadRequestPDU, PlayerFileDownloadResponsePDU
+    PlayerFileDownloadCompletePDU, PlayerFileDownloadRequestPDU, PlayerFileDownloadResponsePDU, PlayerPDU
+from pyrdp.parser import ClientConnectionParser
+from pyrdp.player import LiveTab
 from pyrdp.player.FileDownloadDialog import FileDownloadDialog
 from pyrdp.player.filesystem import DirectoryObserver, Drive, File, FileSystem, FileSystemItemType
 from pyrdp.player.PlayerEventHandler import PlayerEventHandler
@@ -25,18 +28,45 @@ class LiveEventHandler(PlayerEventHandler, DirectoryObserver):
     file read events.
     """
 
-    def __init__(self, viewer: QRemoteDesktop, text: QTextEdit, log: LoggerAdapter, fileSystem: FileSystem, layer: PlayerLayer):
+    connectionClosed = Signal(object)
+    renameTab = Signal(object, str)
+
+    def __init__(self, viewer: QRemoteDesktop, text: QTextEdit, log: LoggerAdapter, fileSystem: FileSystem, layer: PlayerLayer, tabInstance: LiveTab):
         super().__init__(viewer, text)
         self.log = log
         self.fileSystem = fileSystem
         self.layer = layer
         self.drives: Dict[int, Drive] = {}
         self.downloadFiles: Dict[str, BinaryIO] = {}
         self.downloadDialogs: Dict[str, FileDownloadDialog] = {}
+        self.tabInstance = tabInstance
 
         self.handlers[PlayerPDUType.DIRECTORY_LISTING_RESPONSE] = self.handleDirectoryListingResponse
         self.handlers[PlayerPDUType.FILE_DOWNLOAD_RESPONSE] = self.handleDownloadResponse
         self.handlers[PlayerPDUType.FILE_DOWNLOAD_COMPLETE] = self.handleDownloadComplete
+        self.handlers[PlayerPDUType.CLIENT_DATA] = self.onClientData
+        self.handlers[PlayerPDUType.CONNECTION_CLOSE] = self.onConnectionClose
+
+
+    def onClientData(self, pdu: PlayerPDU):
+        """
+        Message the LiveWindow to rename the tab to the hostname of the client
+        """
+
+        clientDataPDU = ClientConnectionParser().parse(pdu.payload)
+        clientName = clientDataPDU.coreData.clientName.strip("\x00")
+
+        self.renameTab.emit(self.tabInstance, clientName)
+        super().onClientData(pdu)
+
+
+    def onConnectionClose(self, pdu: PlayerPDU):
+        """
+        Message the LiveWindow that this tab's connection is closed
+        """
+
+        self.connectionClosed.emit(self.tabInstance)
+        super().onConnectionClose(pdu)
 
     def onDeviceMapping(self, pdu: PlayerDeviceMappingPDU):
         super().onDeviceMapping(pdu)

pyrdp/player/LiveTab.py

@@ -6,7 +6,7 @@
 
 import asyncio
 
-from PySide2.QtCore import Qt, Signal
+from PySide2.QtCore import Qt
 from PySide2.QtWidgets import QHBoxLayout, QWidget
 
 from pyrdp.player.AttackerBar import AttackerBar
@@ -23,8 +23,6 @@ class LiveTab(BaseTab, DirectoryObserver):
     Tab playing a live RDP connection as data is being received over the network.
     """
 
-    connectionClosed = Signal(object)
-
     def __init__(self, parent: QWidget = None):
         layers = AsyncIOPlayerLayerSet()
         rdpWidget = RDPMITMWidget(1024, 768, layers.player)
@@ -33,8 +31,10 @@ def __init__(self, parent: QWidget = None):
         self.layers = layers
         self.rdpWidget = rdpWidget
         self.fileSystem = FileSystem()
-        self.eventHandler = LiveEventHandler(self.widget, self.text, self.log, self.fileSystem, self.layers.player)
+        self.eventHandler = LiveEventHandler(self.widget, self.text, self.log, self.fileSystem, self.layers.player, self)
         self.attackerBar = AttackerBar()
+        self.connectionClosed = self.eventHandler.connectionClosed
+        self.renameTab = self.eventHandler.renameTab
 
         self.attackerBar.controlTaken.connect(lambda: self.rdpWidget.setControlState(True))
         self.attackerBar.controlReleased.connect(lambda: self.rdpWidget.setControlState(False))
@@ -56,9 +56,6 @@ def __init__(self, parent: QWidget = None):
     def getProtocol(self) -> asyncio.Protocol:
         return self.layers.tcp
 
-    def onDisconnection(self):
-        self.connectionClosed.emit()
-
     def onClose(self):
         self.layers.tcp.disconnect(True)
 

pyrdp/player/LiveWindow.py

@@ -6,9 +6,10 @@
 
 import asyncio
 from queue import Queue
+from typing import Dict
 
-from PySide2.QtCore import Signal, Qt
-from PySide2.QtWidgets import QApplication, QWidget
+from PySide2.QtCore import Qt, Signal
+from PySide2.QtWidgets import QApplication, QMessageBox, QWidget
 
 from pyrdp.player.BaseWindow import BaseWindow
 from pyrdp.player.LiveTab import LiveTab
@@ -19,16 +20,20 @@ class LiveWindow(BaseWindow):
     """
     Class that holds logic for live player (network RDP connections as they happen) tabs.
     """
+
     connectionReceived = Signal()
+    closedTabText = " - Closed"
+
+    def __init__(self, address: str, port: int, updateCountSignal: Signal, options: Dict[str, object], parent: QWidget = None):
+        super().__init__(options, parent)
 
-    def __init__(self, address: str, port: int, parent: QWidget = None):
-        super().__init__(parent)
         QApplication.instance().aboutToQuit.connect(self.onClose)
 
         self.server = LiveThread(address, port, self.onConnection)
         self.server.start()
         self.connectionReceived.connect(self.createLivePlayerTab)
         self.queue = Queue()
+        self.updateCountSignal = updateCountSignal
 
     def onConnection(self) -> asyncio.Protocol:
         self.connectionReceived.emit()
@@ -37,19 +42,29 @@ def onConnection(self) -> asyncio.Protocol:
 
     def createLivePlayerTab(self):
         tab = LiveTab()
+        tab.renameTab.connect(self.renameLivePlayerTab)
         tab.connectionClosed.connect(self.onConnectionClosed)
         self.addTab(tab, "New connection")
-        self.setCurrentIndex(self.count() - 1)
+
+        if self.options.get("focusNewTab"):
+            self.setCurrentIndex(self.count() - 1)
+
+        self.updateCountSignal.emit()
         self.queue.put(tab)
 
-    def onConnectionClosed(self, tab: LiveTab):
+    def renameLivePlayerTab(self, tab: LiveTab, name: str):
         index = self.indexOf(tab)
-        text = self.tabText(index)
-        self.setTabText(index, text + " - Closed")
+        self.setTabText(index, name)
 
     def onClose(self):
         self.server.stop()
 
+    def onConnectionClosed(self, tab: LiveTab):
+        index = self.indexOf(tab)
+        text = self.tabText(index)
+        name = text + self.closedTabText
+        self.setTabText(index, name)
+
     def sendKeySequence(self, keys: [Qt.Key]):
         tab: LiveTab = self.currentWidget()
 
@@ -60,4 +75,25 @@ def sendText(self, text: str):
         tab: LiveTab = self.currentWidget()
 
         if tab is not None:
-            tab.sendText(text)
+            tab.sendText(text)
+
+    def onTabClosed(self, index: int):
+        """
+        Gracefully closes the tab by calling the onClose method
+        :param index: Index of the closed tab
+        """
+        super().onTabClosed(index)
+        self.updateCountSignal.emit()
+
+    def onTabCloseRequest(self, index: int):
+        """
+        Prompt the user for validation when the connection is live, then forward call to the parent.
+        """
+        text = self.tabText(index)
+
+        if not text.endswith(self.closedTabText):
+            reply = QMessageBox.question(self, "Confirm close", "Are you sure you want to close a tab with an active connection?", QMessageBox.Yes|QMessageBox.No)
+            if reply == QMessageBox.No:
+                return
+
+        super().onTabCloseRequest(index)

pyrdp/player/MainWindow.py

@@ -4,7 +4,7 @@
 # Licensed under the GPLv3 or later.
 #
 
-from PySide2.QtCore import Qt
+from PySide2.QtCore import Qt, Signal
 from PySide2.QtWidgets import QAction, QFileDialog, QMainWindow, QTabWidget, QInputDialog
 
 from pyrdp.player.LiveWindow import LiveWindow
@@ -16,6 +16,8 @@ class MainWindow(QMainWindow):
     Main window for the player application.
     """
 
+    updateCountSignal = Signal()
+
     def __init__(self, bind_address: str, port: int, filesToRead: [str]):
         """
         :param bind_address: address to bind to when listening for live connections.
@@ -24,19 +26,27 @@ def __init__(self, bind_address: str, port: int, filesToRead: [str]):
         """
         super().__init__()
 
-        self.liveWindow = LiveWindow(bind_address, port)
-        self.replayWindow = ReplayWindow()
+        # TODO : Rework into a class if we add more options later.
+        self.options = {
+            "focusNewTab": True,        # Useful whenever we are getting flooded with connections (or scanned), and we only want to monitor one at a time.
+            "closeTabOnCtrlW": True     # Allow user to toggle Ctrl+W passthrough.
+        }
 
+        self.liveWindow = LiveWindow(bind_address, port, self.updateCountSignal, self.options)
+        self.replayWindow = ReplayWindow(self.options)
         self.tabManager = QTabWidget()
         self.tabManager.addTab(self.liveWindow, "Live connections")
         self.tabManager.addTab(self.replayWindow, "Replays")
         self.setCentralWidget(self.tabManager)
+        self.updateCountSignal.connect(self.updateTabConnectionCount)
 
+        # File menu
         openAction = QAction("Open...", self)
         openAction.setShortcut("Ctrl+O")
         openAction.setStatusTip("Open a replay file")
         openAction.triggered.connect(self.onOpenFile)
 
+        # Command menu
         windowsRAction = QAction("Windows+R", self)
         windowsRAction.setShortcut("Ctrl+Alt+R")
         windowsRAction.setStatusTip("Send a Windows+R key sequence")
@@ -57,6 +67,18 @@ def __init__(self, bind_address: str, port: int, filesToRead: [str]):
         typeTextAction.setStatusTip("Simulate typing on the keyboard")
         typeTextAction.triggered.connect(self.sendText)
 
+        # Options menu
+        focusTabAction = QAction("Focus new connections", self)
+        focusTabAction.setCheckable(True)
+        focusTabAction.setChecked(self.options.get("focusNewTab"))
+        focusTabAction.triggered.connect(lambda: self.toggleFocusNewTab())
+
+        closeTabOnCtrlW = QAction("Close current tab on Ctrl+W", self)
+        closeTabOnCtrlW.setCheckable(True)
+        closeTabOnCtrlW.setChecked(self.options.get("closeTabOnCtrlW"))
+        closeTabOnCtrlW.triggered.connect(lambda: self.toggleCloseTabOnCtrlW())
+
+        # Create menu
         menuBar = self.menuBar()
 
         fileMenu = menuBar.addMenu("File")
@@ -68,6 +90,10 @@ def __init__(self, bind_address: str, port: int, filesToRead: [str]):
         commandMenu.addAction(windowsEAction)
         commandMenu.addAction(typeTextAction)
 
+        optionsMenu = menuBar.addMenu("Options")
+        optionsMenu.addAction(focusTabAction)
+        optionsMenu.addAction(closeTabOnCtrlW)
+
         for fileName in filesToRead:
             self.replayWindow.openFile(fileName)
 
@@ -93,4 +119,17 @@ def sendText(self):
         if not success:
             return
 
-        self.liveWindow.sendText(text)
+        self.liveWindow.sendText(text)
+
+    def toggleFocusNewTab(self):
+        self.options["focusNewTab"] = not self.options.get("focusNewTab")
+
+    def toggleCloseTabOnCtrlW(self):
+        self.options["closeTabOnCtrlW"] = not self.options.get("closeTabOnCtrlW")
+
+    def updateTabConnectionCount(self):
+        """
+        Update the first tab (Live connections) with the current number of tabs
+        """
+
+        self.tabManager.setTabText(0, "Live connections (%d)" % self.liveWindow.count())

pyrdp/player/PlayerEventHandler.py

@@ -6,6 +6,7 @@
 
 from typing import Optional, Union
 
+from PySide2.QtCore import QObject
 from PySide2.QtGui import QTextCursor
 from PySide2.QtWidgets import QTextEdit
 
@@ -22,7 +23,7 @@
 from pyrdp.ui import QRemoteDesktop, RDPBitmapToQtImage
 
 
-class PlayerEventHandler(Observer):
+class PlayerEventHandler(QObject, Observer):
     """
     Class to handle events coming to the player.
     """

pyrdp/player/ReplayWindow.py

@@ -4,8 +4,9 @@
 # Licensed under the GPLv3 or later.
 #
 
-from PySide2.QtGui import QKeySequence
-from PySide2.QtWidgets import QShortcut, QWidget
+from typing import Dict
+
+from PySide2.QtWidgets import QWidget
 
 from pyrdp.player.BaseWindow import BaseWindow
 from pyrdp.player.ReplayTab import ReplayTab
@@ -16,9 +17,8 @@ class ReplayWindow(BaseWindow):
     Class for managing replay tabs.
     """
 
-    def __init__(self, parent: QWidget = None):
-        super().__init__(parent)
-        self.closeTabShortcut = QShortcut(QKeySequence("Ctrl+W"), self, self.closeCurrentTab)
+    def __init__(self, options: Dict[str, object], parent: QWidget = None):
+        super().__init__(options, parent)
 
     def openFile(self, fileName: str):
         """
@@ -28,7 +28,3 @@ def openFile(self, fileName: str):
         tab = ReplayTab(fileName)
         self.addTab(tab, fileName)
         self.log.debug("Loading replay file %(arg1)s", {"arg1": fileName})
-
-    def closeCurrentTab(self):
-        if self.count() > 0:
-            self.onTabClosed(self.currentIndex())