python/qemu/qmp/qmp_tui.py

   1 # Copyright (c) 2021
   2 #
   3 # Authors:
   4 #  Niteesh Babu G S <niteesh.gs@gmail.com>
   5 #
   6 # This work is licensed under the terms of the GNU LGPL, version 2 or
   7 # later.  See the COPYING file in the top-level directory.
   8 """
   9 QMP TUI
  10
  11 QMP TUI is an asynchronous interface built on top the of the QMP library.
  12 It is the successor of QMP-shell and is bought-in as a replacement for it.
  13
  14 Example Usage: qmp-tui <SOCKET | TCP IP:PORT>
  15 Full Usage: qmp-tui --help
  16 """
  17
  18 import argparse
  19 import asyncio
  20 import json
  21 import logging
  22 from logging import Handler, LogRecord
  23 import signal
  24 from typing import (
  25     List,
  26     Optional,
  27     Tuple,
  28     Type,
  29     Union,
  30     cast,
  31 )
  32
  33 from pygments import lexers
  34 from pygments import token as Token
  35 import urwid
  36 import urwid_readline
  37
  38 from .error import ProtocolError
  39 from .legacy import QEMUMonitorProtocol, QMPBadPortError
  40 from .message import DeserializationError, Message, UnexpectedTypeError
  41 from .protocol import ConnectError, Runstate
  42 from .qmp_client import ExecInterruptedError, QMPClient
  43 from .util import create_task, pretty_traceback
  44
  45
  46 # The name of the signal that is used to update the history list
  47 UPDATE_MSG: str = 'UPDATE_MSG'
  48
  49
  50 palette = [
  51     (Token.Punctuation, '', '', '', 'h15,bold', 'g7'),
  52     (Token.Text, '', '', '', '', 'g7'),
  53     (Token.Name.Tag, '', '', '', 'bold,#f88', 'g7'),
  54     (Token.Literal.Number.Integer, '', '', '', '#fa0', 'g7'),
  55     (Token.Literal.String.Double, '', '', '', '#6f6', 'g7'),
  56     (Token.Keyword.Constant, '', '', '', '#6af', 'g7'),
  57     ('DEBUG', '', '', '', '#ddf', 'g7'),
  58     ('INFO', '', '', '', 'g100', 'g7'),
  59     ('WARNING', '', '', '', '#ff6', 'g7'),
  60     ('ERROR', '', '', '', '#a00', 'g7'),
  61     ('CRITICAL', '', '', '', '#a00', 'g7'),
  62     ('background', '', 'black', '', '', 'g7'),
  63 ]
  64
  65
  66 def format_json(msg: str) -> str:
  67     """
  68     Formats valid/invalid multi-line JSON message into a single-line message.
  69
  70     Formatting is first tried using the standard json module. If that fails
  71     due to an decoding error then a simple string manipulation is done to
  72     achieve a single line JSON string.
  73
  74     Converting into single line is more asthetically pleasing when looking
  75     along with error messages.
  76
  77     Eg:
  78     Input:
  79           [ 1,
  80             true,
  81             3 ]
  82     The above input is not a valid QMP message and produces the following error
  83     "QMP message is not a JSON object."
  84     When displaying this in TUI in multiline mode we get
  85
  86         [ 1,
  87           true,
  88           3 ]: QMP message is not a JSON object.
  89
  90     whereas in singleline mode we get the following
  91
  92         [1, true, 3]: QMP message is not a JSON object.
  93
  94     The single line mode is more asthetically pleasing.
  95
  96     :param msg:
  97         The message to formatted into single line.
  98
  99     :return: Formatted singleline message.
 100     """
 101     try:
 102         msg = json.loads(msg)
 103         return str(json.dumps(msg))
 104     except json.decoder.JSONDecodeError:
 105         msg = msg.replace('\n', '')
 106         words = msg.split(' ')
 107         words = list(filter(None, words))
 108         return ' '.join(words)
 109
 110
 111 def has_handler_type(logger: logging.Logger,
 112                      handler_type: Type[Handler]) -> bool:
 113     """
 114     The Logger class has no interface to check if a certain type of handler is
 115     installed or not. So we provide an interface to do so.
 116
 117     :param logger:
 118         Logger object
 119     :param handler_type:
 120         The type of the handler to be checked.
 121
 122     :return: returns True if handler of type `handler_type`.
 123     """
 124     for handler in logger.handlers:
 125         if isinstance(handler, handler_type):
 126             return True
 127     return False
 128
 129
 130 class App(QMPClient):
 131     """
 132     Implements the QMP TUI.
 133
 134     Initializes the widgets and starts the urwid event loop.
 135
 136     :param address:
 137         Address of the server to connect to.
 138     :param num_retries:
 139         The number of times to retry before stopping to reconnect.
 140     :param retry_delay:
 141         The delay(sec) before each retry
 142     """
 143     def __init__(self, address: Union[str, Tuple[str, int]], num_retries: int,
 144                  retry_delay: Optional[int]) -> None:
 145         urwid.register_signal(type(self), UPDATE_MSG)
 146         self.window = Window(self)
 147         self.address = address
 148         self.aloop: Optional[asyncio.AbstractEventLoop] = None
 149         self.num_retries = num_retries
 150         self.retry_delay = retry_delay if retry_delay else 2
 151         self.retry: bool = False
 152         self.exiting: bool = False
 153         super().__init__()
 154
 155     def add_to_history(self, msg: str, level: Optional[str] = None) -> None:
 156         """
 157         Appends the msg to the history list.
 158
 159         :param msg:
 160             The raw message to be appended in string type.
 161         """
 162         urwid.emit_signal(self, UPDATE_MSG, msg, level)
 163
 164     def _cb_outbound(self, msg: Message) -> Message:
 165         """
 166         Callback: outbound message hook.
 167
 168         Appends the outgoing messages to the history box.
 169
 170         :param msg: raw outbound message.
 171         :return: final outbound message.
 172         """
 173         str_msg = str(msg)
 174
 175         if not has_handler_type(logging.getLogger(), TUILogHandler):
 176             logging.debug('Request: %s', str_msg)
 177         self.add_to_history('<-- ' + str_msg)
 178         return msg
 179
 180     def _cb_inbound(self, msg: Message) -> Message:
 181         """
 182         Callback: outbound message hook.
 183
 184         Appends the incoming messages to the history box.
 185
 186         :param msg: raw inbound message.
 187         :return: final inbound message.
 188         """
 189         str_msg = str(msg)
 190
 191         if not has_handler_type(logging.getLogger(), TUILogHandler):
 192             logging.debug('Request: %s', str_msg)
 193         self.add_to_history('--> ' + str_msg)
 194         return msg
 195
 196     async def _send_to_server(self, msg: Message) -> None:
 197         """
 198         This coroutine sends the message to the server.
 199         The message has to be pre-validated.
 200
 201         :param msg:
 202             Pre-validated message to be to sent to the server.
 203
 204         :raise Exception: When an unhandled exception is caught.
 205         """
 206         try:
 207             await self._raw(msg, assign_id='id' not in msg)
 208         except ExecInterruptedError as err:
 209             logging.info('Error server disconnected before reply %s', str(err))
 210             self.add_to_history('Server disconnected before reply', 'ERROR')
 211         except Exception as err:
 212             logging.error('Exception from _send_to_server: %s', str(err))
 213             raise err
 214
 215     def cb_send_to_server(self, raw_msg: str) -> None:
 216         """
 217         Validates and sends the message to the server.
 218         The raw string message is first converted into a Message object
 219         and is then sent to the server.
 220
 221         :param raw_msg:
 222             The raw string message to be sent to the server.
 223
 224         :raise Exception: When an unhandled exception is caught.
 225         """
 226         try:
 227             msg = Message(bytes(raw_msg, encoding='utf-8'))
 228             create_task(self._send_to_server(msg))
 229         except (DeserializationError, UnexpectedTypeError) as err:
 230             raw_msg = format_json(raw_msg)
 231             logging.info('Invalid message: %s', err.error_message)
 232             self.add_to_history(f'{raw_msg}: {err.error_message}', 'ERROR')
 233
 234     def unhandled_input(self, key: str) -> None:
 235         """
 236         Handle's keys which haven't been handled by the child widgets.
 237
 238         :param key:
 239             Unhandled key
 240         """
 241         if key == 'esc':
 242             self.kill_app()
 243
 244     def kill_app(self) -> None:
 245         """
 246         Initiates killing of app. A bridge between asynchronous and synchronous
 247         code.
 248         """
 249         create_task(self._kill_app())
 250
 251     async def _kill_app(self) -> None:
 252         """
 253         This coroutine initiates the actual disconnect process and calls
 254         urwid.ExitMainLoop() to kill the TUI.
 255
 256         :raise Exception: When an unhandled exception is caught.
 257         """
 258         self.exiting = True
 259         await self.disconnect()
 260         logging.debug('Disconnect finished. Exiting app')
 261         raise urwid.ExitMainLoop()
 262
 263     async def disconnect(self) -> None:
 264         """
 265         Overrides the disconnect method to handle the errors locally.
 266         """
 267         try:
 268             await super().disconnect()
 269         except (OSError, EOFError) as err:
 270             logging.info('disconnect: %s', str(err))
 271             self.retry = True
 272         except ProtocolError as err:
 273             logging.info('disconnect: %s', str(err))
 274         except Exception as err:
 275             logging.error('disconnect: Unhandled exception %s', str(err))
 276             raise err
 277
 278     def _set_status(self, msg: str) -> None:
 279         """
 280         Sets the message as the status.
 281
 282         :param msg:
 283             The message to be displayed in the status bar.
 284         """
 285         self.window.footer.set_text(msg)
 286
 287     def _get_formatted_address(self) -> str:
 288         """
 289         Returns a formatted version of the server's address.
 290
 291         :return: formatted address
 292         """
 293         if isinstance(self.address, tuple):
 294             host, port = self.address
 295             addr = f'{host}:{port}'
 296         else:
 297             addr = f'{self.address}'
 298         return addr
 299
 300     async def _initiate_connection(self) -> Optional[ConnectError]:
 301         """
 302         Tries connecting to a server a number of times with a delay between
 303         each try. If all retries failed then return the error faced during
 304         the last retry.
 305
 306         :return: Error faced during last retry.
 307         """
 308         current_retries = 0
 309         err = None
 310
 311         # initial try
 312         await self.connect_server()
 313         while self.retry and current_retries < self.num_retries:
 314             logging.info('Connection Failed, retrying in %d', self.retry_delay)
 315             status = f'[Retry #{current_retries} ({self.retry_delay}s)]'
 316             self._set_status(status)
 317
 318             await asyncio.sleep(self.retry_delay)
 319
 320             err = await self.connect_server()
 321             current_retries += 1
 322         # If all retries failed report the last error
 323         if err:
 324             logging.info('All retries failed: %s', err)
 325             return err
 326         return None
 327
 328     async def manage_connection(self) -> None:
 329         """
 330         Manage the connection based on the current run state.
 331
 332         A reconnect is issued when the current state is IDLE and the number
 333         of retries is not exhausted.
 334         A disconnect is issued when the current state is DISCONNECTING.
 335         """
 336         while not self.exiting:
 337             if self.runstate == Runstate.IDLE:
 338                 err = await self._initiate_connection()
 339                 # If retry is still true then, we have exhausted all our tries.
 340                 if err:
 341                     self._set_status(f'[Error: {err.error_message}]')
 342                 else:
 343                     addr = self._get_formatted_address()
 344                     self._set_status(f'[Connected {addr}]')
 345             elif self.runstate == Runstate.DISCONNECTING:
 346                 self._set_status('[Disconnected]')
 347                 await self.disconnect()
 348                 # check if a retry is needed
 349                 if self.runstate == Runstate.IDLE:
 350                     continue
 351             await self.runstate_changed()
 352
 353     async def connect_server(self) -> Optional[ConnectError]:
 354         """
 355         Initiates a connection to the server at address `self.address`
 356         and in case of a failure, sets the status to the respective error.
 357         """
 358         try:
 359             await self.connect(self.address)
 360             self.retry = False
 361         except ConnectError as err:
 362             logging.info('connect_server: ConnectError %s', str(err))
 363             self.retry = True
 364             return err
 365         return None
 366
 367     def run(self, debug: bool = False) -> None:
 368         """
 369         Starts the long running co-routines and the urwid event loop.
 370
 371         :param debug:
 372             Enables/Disables asyncio event loop debugging
 373         """
 374         screen = urwid.raw_display.Screen()
 375         screen.set_terminal_properties(256)
 376
 377         self.aloop = asyncio.get_event_loop()
 378         self.aloop.set_debug(debug)
 379
 380         # Gracefully handle SIGTERM and SIGINT signals
 381         cancel_signals = [signal.SIGTERM, signal.SIGINT]
 382         for sig in cancel_signals:
 383             self.aloop.add_signal_handler(sig, self.kill_app)
 384
 385         event_loop = urwid.AsyncioEventLoop(loop=self.aloop)
 386         main_loop = urwid.MainLoop(urwid.AttrMap(self.window, 'background'),
 387                                    unhandled_input=self.unhandled_input,
 388                                    screen=screen,
 389                                    palette=palette,
 390                                    handle_mouse=True,
 391                                    event_loop=event_loop)
 392
 393         create_task(self.manage_connection(), self.aloop)
 394         try:
 395             main_loop.run()
 396         except Exception as err:
 397             logging.error('%s\n%s\n', str(err), pretty_traceback())
 398             raise err
 399
 400
 401 class StatusBar(urwid.Text):
 402     """
 403     A simple statusbar modelled using the Text widget. The status can be
 404     set using the set_text function. All text set is aligned to right.
 405
 406     :param text: Initial text to be displayed. Default is empty str.
 407     """
 408     def __init__(self, text: str = ''):
 409         super().__init__(text, align='right')
 410
 411
 412 class Editor(urwid_readline.ReadlineEdit):
 413     """
 414     A simple editor modelled using the urwid_readline.ReadlineEdit widget.
 415     Mimcs GNU readline shortcuts and provides history support.
 416
 417     The readline shortcuts can be found below:
 418     https://github.com/rr-/urwid_readline#features
 419
 420     Along with the readline features, this editor also has support for
 421     history. Pressing the 'up'/'down' switches between the prev/next messages
 422     available in the history.
 423
 424     Currently there is no support to save the history to a file. The history of
 425     previous commands is lost on exit.
 426
 427     :param parent: Reference to the TUI object.
 428     """
 429     def __init__(self, parent: App) -> None:
 430         super().__init__(caption='> ', multiline=True)
 431         self.parent = parent
 432         self.history: List[str] = []
 433         self.last_index: int = -1
 434         self.show_history: bool = False
 435
 436     def keypress(self, size: Tuple[int, int], key: str) -> Optional[str]:
 437         """
 438         Handles the keypress on this widget.
 439
 440         :param size:
 441             The current size of the widget.
 442         :param key:
 443             The key to be handled.
 444
 445         :return: Unhandled key if any.
 446         """
 447         msg = self.get_edit_text()
 448         if key == 'up' and not msg:
 449             # Show the history when 'up arrow' is pressed with no input text.
 450             # NOTE: The show_history logic is necessary because in 'multiline'
 451             # mode (which we use) 'up arrow' is used to move between lines.
 452             if not self.history:
 453                 return None
 454             self.show_history = True
 455             last_msg = self.history[self.last_index]
 456             self.set_edit_text(last_msg)
 457             self.edit_pos = len(last_msg)
 458         elif key == 'up' and self.show_history:
 459             self.last_index = max(self.last_index - 1, -len(self.history))
 460             self.set_edit_text(self.history[self.last_index])
 461             self.edit_pos = len(self.history[self.last_index])
 462         elif key == 'down' and self.show_history:
 463             if self.last_index == -1:
 464                 self.set_edit_text('')
 465                 self.show_history = False
 466             else:
 467                 self.last_index += 1
 468                 self.set_edit_text(self.history[self.last_index])
 469                 self.edit_pos = len(self.history[self.last_index])
 470         elif key == 'meta enter':
 471             # When using multiline, enter inserts a new line into the editor
 472             # send the input to the server on alt + enter
 473             self.parent.cb_send_to_server(msg)
 474             self.history.append(msg)
 475             self.set_edit_text('')
 476             self.last_index = -1
 477             self.show_history = False
 478         else:
 479             self.show_history = False
 480             self.last_index = -1
 481             return cast(Optional[str], super().keypress(size, key))
 482         return None
 483
 484
 485 class EditorWidget(urwid.Filler):
 486     """
 487     Wrapper around the editor widget.
 488
 489     The Editor is a flow widget and has to wrapped inside a box widget.
 490     This class wraps the Editor inside filler widget.
 491
 492     :param parent: Reference to the TUI object.
 493     """
 494     def __init__(self, parent: App) -> None:
 495         super().__init__(Editor(parent), valign='top')
 496
 497
 498 class HistoryBox(urwid.ListBox):
 499     """
 500     This widget is modelled using the ListBox widget, contains the list of
 501     all messages both QMP messages and log messsages to be shown in the TUI.
 502
 503     The messages are urwid.Text widgets. On every append of a message, the
 504     focus is shifted to the last appended message.
 505
 506     :param parent: Reference to the TUI object.
 507     """
 508     def __init__(self, parent: App) -> None:
 509         self.parent = parent
 510         self.history = urwid.SimpleFocusListWalker([])
 511         super().__init__(self.history)
 512
 513     def add_to_history(self,
 514                        history: Union[str, List[Tuple[str, str]]]) -> None:
 515         """
 516         Appends a message to the list and set the focus to the last appended
 517         message.
 518
 519         :param history:
 520             The history item(message/event) to be appended to the list.
 521         """
 522         self.history.append(urwid.Text(history))
 523         self.history.set_focus(len(self.history) - 1)
 524
 525     def mouse_event(self, size: Tuple[int, int], _event: str, button: float,
 526                     _x: int, _y: int, focus: bool) -> None:
 527         # Unfortunately there are no urwid constants that represent the mouse
 528         # events.
 529         if button == 4:  # Scroll up event
 530             super().keypress(size, 'up')
 531         elif button == 5:  # Scroll down event
 532             super().keypress(size, 'down')
 533
 534
 535 class HistoryWindow(urwid.Frame):
 536     """
 537     This window composes the HistoryBox and EditorWidget in a horizontal split.
 538     By default the first focus is given to the history box.
 539
 540     :param parent: Reference to the TUI object.
 541     """
 542     def __init__(self, parent: App) -> None:
 543         self.parent = parent
 544         self.editor_widget = EditorWidget(parent)
 545         self.editor = urwid.LineBox(self.editor_widget)
 546         self.history = HistoryBox(parent)
 547         self.body = urwid.Pile([('weight', 80, self.history),
 548                                 ('weight', 20, self.editor)])
 549         super().__init__(self.body)
 550         urwid.connect_signal(self.parent, UPDATE_MSG, self.cb_add_to_history)
 551
 552     def cb_add_to_history(self, msg: str, level: Optional[str] = None) -> None:
 553         """
 554         Appends a message to the history box
 555
 556         :param msg:
 557             The message to be appended to the history box.
 558         :param level:
 559             The log level of the message, if it is a log message.
 560         """
 561         formatted = []
 562         if level:
 563             msg = f'[{level}]: {msg}'
 564             formatted.append((level, msg))
 565         else:
 566             lexer = lexers.JsonLexer()  # pylint: disable=no-member
 567             for token in lexer.get_tokens(msg):
 568                 formatted.append(token)
 569         self.history.add_to_history(formatted)
 570
 571
 572 class Window(urwid.Frame):
 573     """
 574     This window is the top most widget of the TUI and will contain other
 575     windows. Each child of this widget is responsible for displaying a specific
 576     functionality.
 577
 578     :param parent: Reference to the TUI object.
 579     """
 580     def __init__(self, parent: App) -> None:
 581         self.parent = parent
 582         footer = StatusBar()
 583         body = HistoryWindow(parent)
 584         super().__init__(body, footer=footer)
 585
 586
 587 class TUILogHandler(Handler):
 588     """
 589     This handler routes all the log messages to the TUI screen.
 590     It is installed to the root logger to so that the log message from all
 591     libraries begin used is routed to the screen.
 592
 593     :param tui: Reference to the TUI object.
 594     """
 595     def __init__(self, tui: App) -> None:
 596         super().__init__()
 597         self.tui = tui
 598
 599     def emit(self, record: LogRecord) -> None:
 600         """
 601         Emits a record to the TUI screen.
 602
 603         Appends the log message to the TUI screen
 604         """
 605         level = record.levelname
 606         msg = record.getMessage()
 607         self.tui.add_to_history(msg, level)
 608
 609
 610 def main() -> None:
 611     """
 612     Driver of the whole script, parses arguments, initialize the TUI and
 613     the logger.
 614     """
 615     parser = argparse.ArgumentParser(description='QMP TUI')
 616     parser.add_argument('qmp_server', help='Address of the QMP server. '
 617                         'Format <UNIX socket path | TCP addr:port>')
 618     parser.add_argument('--num-retries', type=int, default=10,
 619                         help='Number of times to reconnect before giving up.')
 620     parser.add_argument('--retry-delay', type=int,
 621                         help='Time(s) to wait before next retry. '
 622                         'Default action is to wait 2s between each retry.')
 623     parser.add_argument('--log-file', help='The Log file name')
 624     parser.add_argument('--log-level', default='WARNING',
 625                         help='Log level <CRITICAL|ERROR|WARNING|INFO|DEBUG|>')
 626     parser.add_argument('--asyncio-debug', action='store_true',
 627                         help='Enable debug mode for asyncio loop. '
 628                         'Generates lot of output, makes TUI unusable when '
 629                         'logs are logged in the TUI. '
 630                         'Use only when logging to a file.')
 631     args = parser.parse_args()
 632
 633     try:
 634         address = QEMUMonitorProtocol.parse_address(args.qmp_server)
 635     except QMPBadPortError as err:
 636         parser.error(str(err))
 637
 638     app = App(address, args.num_retries, args.retry_delay)
 639
 640     root_logger = logging.getLogger()
 641     root_logger.setLevel(logging.getLevelName(args.log_level))
 642
 643     if args.log_file:
 644         root_logger.addHandler(logging.FileHandler(args.log_file))
 645     else:
 646         root_logger.addHandler(TUILogHandler(app))
 647
 648     app.run(args.asyncio_debug)
 649
 650
 651 if __name__ == '__main__':
 652     main()