Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Pyrdp-player improvements #119

Merged
merged 5 commits into from
Jul 29, 2019
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 20 additions & 4 deletions pyrdp/player/BaseWindow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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())
34 changes: 32 additions & 2 deletions pyrdp/player/LiveEventHandler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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")
Res260 marked this conversation as resolved.
Show resolved Hide resolved

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)
Expand Down
11 changes: 4 additions & 7 deletions pyrdp/player/LiveTab.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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))
Expand All @@ -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)

Expand Down
54 changes: 45 additions & 9 deletions pyrdp/player/LiveWindow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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()
Expand All @@ -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()

Expand All @@ -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)
47 changes: 43 additions & 4 deletions pyrdp/player/MainWindow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand All @@ -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")
Expand All @@ -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")
Expand All @@ -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)

Expand All @@ -91,4 +117,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())
3 changes: 2 additions & 1 deletion pyrdp/player/PlayerEventHandler.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

from typing import Optional, Union

from PySide2.QtCore import QObject
from PySide2.QtGui import QTextCursor
from PySide2.QtWidgets import QTextEdit

Expand All @@ -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.
"""
Expand Down
14 changes: 5 additions & 9 deletions pyrdp/player/ReplayWindow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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):
"""
Expand All @@ -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())