mtda-ui

#!/usr/bin/env python3
# ---------------------------------------------------------------------------
# Command Line Interface for MTDA
# ---------------------------------------------------------------------------
#
# This software is a part of MTDA.
# Copyright (C) 2024 Siemens Digital Industries Software
#
# ---------------------------------------------------------------------------
# SPDX-License-Identifier: MIT
# ---------------------------------------------------------------------------

import os
import socket
import sys
import threading
import time

from py3qterm import TerminalWidget
from py3qterm.backend import Session
from py3qterm.backend import Terminal

from PyQt5.Qt import Qt
from PyQt5.QtCore import QUrl

from PyQt5.QtWidgets import (QMainWindow, QAction, QApplication,
                             QHBoxLayout, QListWidget, QListWidgetItem,
                             QSizePolicy, QStyle, QTabWidget, QWidget)

from PyQt5.QtMultimedia import QMediaContent, QMediaPlayer
from PyQt5.QtMultimediaWidgets import QVideoWidget

from mtda.client import Client
from mtda.console.screen import ScreenOutput
import mtda.constants as CONSTS
import mtda.discovery


class DeviceConsole(TerminalWidget):

    def __init__(self, parent):
        super().__init__(parent)

    def agent(self):
        if Session._mux is not None:
            return Session._mux.agent(self._session._session_id)
        return None

    def execute(self, remote=None):
        if remote is None:
            remote = os.getenv("MTDA_REMOTE", None)
        if self._session is not None:
            self._session.stop()
        self._session = DeviceSession()
        if remote is not None:
            self._session.start(remote)
        self._timer_id = None
        self.focusInEvent(None)

    def mouseDoubleClickEvent(self, event):
        pass

    def mouseMoveEvent(self, event):
        pass

    def mousePressEvent(self, event):
        self.focusInEvent(None)

    def mouseReleaseEvent(self, event):
        pass

    def screen(self):
        if Session._mux is not None:
            return Session._mux.screen(self._session._session_id)
        return None


class DeviceMux:

    def __init__(self):
        self.session = {}

    def agent(self, sid):
        if sid in self.session:
            return self.session[sid]['agent']
        return None

    def proc_bury(self, sid):
        if sid not in self.session:
            return False

        agent = self.agent(sid)
        # Detach remote consoles
        agent.console_remote(None, None)
        agent.monitor_remote(None, None)
        # Disconnect from agent
        agent.stop()

        # Session was buried
        return True

    def proc_keepalive(self, sid, w, h, cmd=None):
        if sid not in self.session:
            # Start a new session
            self.remote = cmd
            agent = Client(self.remote)
            term = Terminal(w, h)
            self.session[sid] = {
                'agent': agent,
                'state': 'unborn',
                'term': term,
                'time': time.time(),
                'w': w,
                'h': h
            }
            screen = DeviceScreen(agent, term, self.session[sid])
            agent.console_remote(self.remote, screen)
            agent.monitor_remote(self.remote, screen)
            agent.start()
            self.session[sid]['state'] = 'alive'
            self.session[sid]['screen'] = screen
        os.environ["COLUMNS"] = str(w)
        os.environ["ROWS"] = str(h)
        return True

    def proc_dump(self, sid):
        if sid not in self.session:
            return False
        return self.session[sid]['term'].dump()

    def proc_write(self, sid, d):
        if sid not in self.session:
            return False
        elif self.session[sid]['state'] != 'alive':
            return False
        term = self.session[sid]['term']
        d = term.pipe(d)
        agent = self.session[sid]['agent']
        agent.console_send(d)
        return True

    def screen(self, sid):
        if sid in self.session:
            return self.session[sid]['screen']
        return None

    def stop(self):
        return None


class DeviceScreen(ScreenOutput):

    def __init__(self, mtda, term, session):
        super().__init__(mtda)
        self._event_handler = None
        self._term = term
        self._session = session

    def add_listener(self, handler):
        self._event_handler = handler

    def on_event(self, event):
        if self._event_handler is not None:
            event = event.split()
            self._event_handler(event[0], event[1:])

    def write(self, data):
        self._term.write(data)
        data = self._term.read()
        self._session["changed"] = time.time()


class DeviceSession(Session):

    def __init__(self):
        super().__init__()
        Session._mux = DeviceMux()


class MultiTenantDeviceAccessWindow(QMainWindow):

    def __init__(self):
        super().__init__()

        self.initUI()
        self.service_watcher = mtda.discovery.Watcher(
                CONSTS.MDNS.TYPE,
                self.service_event)
        self.service_thread = threading.Thread(
                target=self.service_watcher.listen,
                daemon=True,
                name='service discovery')
        self.service_thread.start()

    def capture(self):
        screen = self.console.screen()
        if screen is not None:
            if screen.capture_enabled() is True:
                self.capture_stop()
            else:
                self.capture_start()

    def capture_reset(self):
        agent = self.console.agent()
        enabled = (agent is not None)
        self.captureAct.setEnabled(enabled)

    def capture_start(self):
        screen = self.console.screen()
        if screen is not None:
            if screen.capture_enabled() is False:
                screen.capture_start()
            icn = self.style().standardIcon(QStyle.SP_BrowserStop)
            self.captureAct.setIcon(icn)

    def capture_stop(self):
        screen = self.console.screen()
        if screen is not None:
            if screen.capture_enabled() is True:
                screen.capture_stop()
            icn = self.style().standardIcon(QStyle.SP_DialogSaveButton)
            self.captureAct.setIcon(icn)

    def device_event(self, event, data):
        if event == "POWER":
            self.power_event(data[0])
        elif event == "STORAGE":
            self.storage_event(data[0])

    def exit(self):
        self.service_watcher.shutdown()
        self.close()

    def initUI(self):

        self.root_widget = QWidget()
        self.hlayout = QHBoxLayout()
        self.root_widget.setLayout(self.hlayout)
        self.setCentralWidget(self.root_widget)

        self.services = QListWidget()
        self.services.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Minimum)
        self.services.itemDoubleClicked.connect(self.service_change)
        self.hlayout.addWidget(self.services)

        self.tabs = QTabWidget(self.root_widget)
        self.tabs.setTabBarAutoHide(True)
        self.hlayout.addWidget(self.tabs)

        self.console = DeviceConsole(self.tabs)
        self.tabs.addTab(self.console, 'Console')

        self.media = QMediaPlayer(None, QMediaPlayer.StreamPlayback)
        self.video = QVideoWidget(self.tabs)
        self.media.setVideoOutput(self.video)

        self.video_reset()

        exitIcn = self.style().standardIcon(QStyle.SP_FileDialogEnd)
        exitAct = QAction(exitIcn, 'Exit', self)
        exitAct.setShortcut('Ctrl+Q')
        exitAct.setStatusTip('Exit MTDA')
        exitAct.triggered.connect(self.exit)

        powerAct = QAction('Power on/off', self)
        powerAct.setShortcut('Ctrl+P')
        powerAct.triggered.connect(self.power)
        self.powerAct = powerAct
        self.power_reset()

        storageAct = QAction('Swap storage', self)
        storageAct.setShortcut('Ctrl+S')
        storageAct.triggered.connect(self.storage)
        self.storageAct = storageAct
        self.storage_reset()

        captureIcn = self.style().standardIcon(QStyle.SP_DialogSaveButton)
        captureAct = QAction(captureIcn, 'Capture screen start/stop', self)
        captureAct.setStatusTip('Capture screen output start/stop')
        captureAct.triggered.connect(self.capture)
        self.captureAct = captureAct
        self.capture_reset()

        self.statusBar()

        menubar = self.menuBar()
        fileMenu = menubar.addMenu('&File')
        fileMenu.addAction(exitAct)

        toolbar = self.addToolBar('Exit')
        toolbar.addAction(powerAct)
        toolbar.addAction(storageAct)
        toolbar.addAction(captureAct)

        self.resize(880, 760)
        self.setWindowTitle('MTDA')
        self.show()

    def power(self):
        agent = self.console.agent()
        if agent is not None:
            agent.target_toggle()

    def power_event(self, event):
        a = self.powerAct
        if event == "OFF":
            icn = self.style().standardIcon(QStyle.SP_ComputerIcon)
            txt = "Power on device"
            self.video_stop()
        else:
            icn = self.style().standardIcon(QStyle.SP_MediaStop)
            txt = "Power off device"
            self.video_start()

        a.setStatusTip(txt)
        a.setText(txt)
        a.setIcon(icn)

    def power_reset(self):
        agent = self.console.agent()
        enabled = (agent is not None)

        # Sync UI state with agent
        self.powerAct.setEnabled(enabled)
        if enabled is True:
            self.power_event(agent.target_status())

    def service_change(self, item):
        info = item.data(Qt.UserRole)
        ip = socket.inet_ntoa(info.addresses[0])

        # Stop on-going screen capture / video playback
        self.capture_stop()
        self.video_stop()

        # Connect to selected remote
        self.video_url = ""
        self.console.execute(ip)

        # Listen to screen events from the agent
        screen = self.console.screen()
        screen.add_listener(self.device_event)

        # Reset widgets
        self.power_reset()
        self.capture_reset()
        self.storage_reset()
        self.video_reset()

    def service_event(self, event, name, data=None):
        if event == "ADD":
            icon = self.style().standardIcon(QStyle.SP_ComputerIcon)
            newService = QListWidgetItem()
            shortname = name.split('.')[0]
            newService.setIcon(icon)
            newService.setText(shortname)
            newService.setData(Qt.ToolTipRole, name)
            newService.setData(Qt.UserRole, data)
            self.services.addItem(newService)
        else:
            count = self.services.count()
            for i in range(0, count):
                item = self.services.item(i)
                if item.data(Qt.ToolTipRole) == name:
                    if event == "REMOVE":
                        self.services.takeItem(i)
                    if event == "UPDATE":
                        item.setData(Qt.UserRole, data)
                    break

    def storage(self):
        agent = self.console.agent()
        if agent is not None:
            agent.storage_swap()

    def storage_event(self, event):
        a = self.storageAct
        if event == "HOST":
            icn = self.style().standardIcon(QStyle.SP_DriveHDIcon)
            txt = "Connect storage to target"
        elif event == "TARGET":
            icn = self.style().standardIcon(QStyle.SP_DirClosedIcon)
            txt = "Connect storage to host"
        else:
            return None
        a.setStatusTip(txt)
        a.setText(txt)
        a.setIcon(icn)

    def storage_reset(self):
        agent = self.console.agent()
        enabled = (agent is not None)
        self.storageAct.setEnabled(enabled)
        if agent is not None:
            status, writing, written = agent.storage_status()
            self.storage_event(status)

    def video_reset(self):
        agent = self.console.agent()
        enabled = (agent is not None)

        self.video_url = None
        if enabled is True:
            opts = {'sink': {'name': 'qtvideosink'}}
            self.video_url = agent.video_url(opts=opts)
        idx = self.tabs.indexOf(self.video)
        if self.video_url is not None:
            # Use provided video url
            self.media.setMedia(QMediaContent(QUrl(self.video_url)))
            if idx < 0:
                self.tabs.addTab(self.video, 'Video')
        elif idx >= 0:
            # Remove Video tab
            self.tabs.removeTab(idx)

    def video_start(self):
        if self.video_url != "":
            self.media.play()

    def video_stop(self):
        self.media.stop()


def main():
    app = QApplication(sys.argv)
    MultiTenantDeviceAccessWindow()
    os._exit(app.exec_())


if __name__ == '__main__':
    main()