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""" 9QMP TUI 10 11QMP TUI is an asynchronous interface built on top the of the QMP library. 12It is the successor of QMP-shell and is bought-in as a replacement for it. 13 14Example Usage: qmp-tui <SOCKET | TCP IP:PORT> 15Full Usage: qmp-tui --help 16""" 17 18import argparse 19import asyncio 20import json 21import logging 22from logging import Handler, LogRecord 23import signal 24from typing import ( 25 List, 26 Optional, 27 Tuple, 28 Type, 29 Union, 30 cast, 31) 32 33from pygments import lexers 34from pygments import token as Token 35import urwid 36import urwid_readline 37 38from .error import ProtocolError 39from .legacy import QEMUMonitorProtocol, QMPBadPortError 40from .message import DeserializationError, Message, UnexpectedTypeError 41from .protocol import ConnectError, Runstate 42from .qmp_client import ExecInterruptedError, QMPClient 43from .util import create_task, pretty_traceback 44 45 46# The name of the signal that is used to update the history list 47UPDATE_MSG: str = 'UPDATE_MSG' 48 49 50palette = [ 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 66def 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 aesthetically 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 aesthetically 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 111def 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 130class 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 # mypy 1.4.0 doesn't believe runstate can change after 350 # disconnect(), hence the cast. 351 state = cast(Runstate, self.runstate) 352 if state == Runstate.IDLE: 353 continue 354 await self.runstate_changed() 355 356 async def connect_server(self) -> Optional[ConnectError]: 357 """ 358 Initiates a connection to the server at address `self.address` 359 and in case of a failure, sets the status to the respective error. 360 """ 361 try: 362 await self.connect(self.address) 363 self.retry = False 364 except ConnectError as err: 365 logging.info('connect_server: ConnectError %s', str(err)) 366 self.retry = True 367 return err 368 return None 369 370 def run(self, debug: bool = False) -> None: 371 """ 372 Starts the long running co-routines and the urwid event loop. 373 374 :param debug: 375 Enables/Disables asyncio event loop debugging 376 """ 377 screen = urwid.raw_display.Screen() 378 screen.set_terminal_properties(256) 379 380 self.aloop = asyncio.get_event_loop() 381 self.aloop.set_debug(debug) 382 383 # Gracefully handle SIGTERM and SIGINT signals 384 cancel_signals = [signal.SIGTERM, signal.SIGINT] 385 for sig in cancel_signals: 386 self.aloop.add_signal_handler(sig, self.kill_app) 387 388 event_loop = urwid.AsyncioEventLoop(loop=self.aloop) 389 main_loop = urwid.MainLoop(urwid.AttrMap(self.window, 'background'), 390 unhandled_input=self.unhandled_input, 391 screen=screen, 392 palette=palette, 393 handle_mouse=True, 394 event_loop=event_loop) 395 396 create_task(self.manage_connection(), self.aloop) 397 try: 398 main_loop.run() 399 except Exception as err: 400 logging.error('%s\n%s\n', str(err), pretty_traceback()) 401 raise err 402 403 404class StatusBar(urwid.Text): 405 """ 406 A simple statusbar modelled using the Text widget. The status can be 407 set using the set_text function. All text set is aligned to right. 408 409 :param text: Initial text to be displayed. Default is empty str. 410 """ 411 def __init__(self, text: str = ''): 412 super().__init__(text, align='right') 413 414 415class Editor(urwid_readline.ReadlineEdit): 416 """ 417 A simple editor modelled using the urwid_readline.ReadlineEdit widget. 418 Mimcs GNU readline shortcuts and provides history support. 419 420 The readline shortcuts can be found below: 421 https://github.com/rr-/urwid_readline#features 422 423 Along with the readline features, this editor also has support for 424 history. Pressing the 'up'/'down' switches between the prev/next messages 425 available in the history. 426 427 Currently there is no support to save the history to a file. The history of 428 previous commands is lost on exit. 429 430 :param parent: Reference to the TUI object. 431 """ 432 def __init__(self, parent: App) -> None: 433 super().__init__(caption='> ', multiline=True) 434 self.parent = parent 435 self.history: List[str] = [] 436 self.last_index: int = -1 437 self.show_history: bool = False 438 439 def keypress(self, size: Tuple[int, int], key: str) -> Optional[str]: 440 """ 441 Handles the keypress on this widget. 442 443 :param size: 444 The current size of the widget. 445 :param key: 446 The key to be handled. 447 448 :return: Unhandled key if any. 449 """ 450 msg = self.get_edit_text() 451 if key == 'up' and not msg: 452 # Show the history when 'up arrow' is pressed with no input text. 453 # NOTE: The show_history logic is necessary because in 'multiline' 454 # mode (which we use) 'up arrow' is used to move between lines. 455 if not self.history: 456 return None 457 self.show_history = True 458 last_msg = self.history[self.last_index] 459 self.set_edit_text(last_msg) 460 self.edit_pos = len(last_msg) 461 elif key == 'up' and self.show_history: 462 self.last_index = max(self.last_index - 1, -len(self.history)) 463 self.set_edit_text(self.history[self.last_index]) 464 self.edit_pos = len(self.history[self.last_index]) 465 elif key == 'down' and self.show_history: 466 if self.last_index == -1: 467 self.set_edit_text('') 468 self.show_history = False 469 else: 470 self.last_index += 1 471 self.set_edit_text(self.history[self.last_index]) 472 self.edit_pos = len(self.history[self.last_index]) 473 elif key == 'meta enter': 474 # When using multiline, enter inserts a new line into the editor 475 # send the input to the server on alt + enter 476 self.parent.cb_send_to_server(msg) 477 self.history.append(msg) 478 self.set_edit_text('') 479 self.last_index = -1 480 self.show_history = False 481 else: 482 self.show_history = False 483 self.last_index = -1 484 return cast(Optional[str], super().keypress(size, key)) 485 return None 486 487 488class EditorWidget(urwid.Filler): 489 """ 490 Wrapper around the editor widget. 491 492 The Editor is a flow widget and has to wrapped inside a box widget. 493 This class wraps the Editor inside filler widget. 494 495 :param parent: Reference to the TUI object. 496 """ 497 def __init__(self, parent: App) -> None: 498 super().__init__(Editor(parent), valign='top') 499 500 501class HistoryBox(urwid.ListBox): 502 """ 503 This widget is modelled using the ListBox widget, contains the list of 504 all messages both QMP messages and log messages to be shown in the TUI. 505 506 The messages are urwid.Text widgets. On every append of a message, the 507 focus is shifted to the last appended message. 508 509 :param parent: Reference to the TUI object. 510 """ 511 def __init__(self, parent: App) -> None: 512 self.parent = parent 513 self.history = urwid.SimpleFocusListWalker([]) 514 super().__init__(self.history) 515 516 def add_to_history(self, 517 history: Union[str, List[Tuple[str, str]]]) -> None: 518 """ 519 Appends a message to the list and set the focus to the last appended 520 message. 521 522 :param history: 523 The history item(message/event) to be appended to the list. 524 """ 525 self.history.append(urwid.Text(history)) 526 self.history.set_focus(len(self.history) - 1) 527 528 def mouse_event(self, size: Tuple[int, int], _event: str, button: float, 529 _x: int, _y: int, focus: bool) -> None: 530 # Unfortunately there are no urwid constants that represent the mouse 531 # events. 532 if button == 4: # Scroll up event 533 super().keypress(size, 'up') 534 elif button == 5: # Scroll down event 535 super().keypress(size, 'down') 536 537 538class HistoryWindow(urwid.Frame): 539 """ 540 This window composes the HistoryBox and EditorWidget in a horizontal split. 541 By default the first focus is given to the history box. 542 543 :param parent: Reference to the TUI object. 544 """ 545 def __init__(self, parent: App) -> None: 546 self.parent = parent 547 self.editor_widget = EditorWidget(parent) 548 self.editor = urwid.LineBox(self.editor_widget) 549 self.history = HistoryBox(parent) 550 self.body = urwid.Pile([('weight', 80, self.history), 551 ('weight', 20, self.editor)]) 552 super().__init__(self.body) 553 urwid.connect_signal(self.parent, UPDATE_MSG, self.cb_add_to_history) 554 555 def cb_add_to_history(self, msg: str, level: Optional[str] = None) -> None: 556 """ 557 Appends a message to the history box 558 559 :param msg: 560 The message to be appended to the history box. 561 :param level: 562 The log level of the message, if it is a log message. 563 """ 564 formatted = [] 565 if level: 566 msg = f'[{level}]: {msg}' 567 formatted.append((level, msg)) 568 else: 569 lexer = lexers.JsonLexer() # pylint: disable=no-member 570 for token in lexer.get_tokens(msg): 571 formatted.append(token) 572 self.history.add_to_history(formatted) 573 574 575class Window(urwid.Frame): 576 """ 577 This window is the top most widget of the TUI and will contain other 578 windows. Each child of this widget is responsible for displaying a specific 579 functionality. 580 581 :param parent: Reference to the TUI object. 582 """ 583 def __init__(self, parent: App) -> None: 584 self.parent = parent 585 footer = StatusBar() 586 body = HistoryWindow(parent) 587 super().__init__(body, footer=footer) 588 589 590class TUILogHandler(Handler): 591 """ 592 This handler routes all the log messages to the TUI screen. 593 It is installed to the root logger to so that the log message from all 594 libraries begin used is routed to the screen. 595 596 :param tui: Reference to the TUI object. 597 """ 598 def __init__(self, tui: App) -> None: 599 super().__init__() 600 self.tui = tui 601 602 def emit(self, record: LogRecord) -> None: 603 """ 604 Emits a record to the TUI screen. 605 606 Appends the log message to the TUI screen 607 """ 608 level = record.levelname 609 msg = record.getMessage() 610 self.tui.add_to_history(msg, level) 611 612 613def main() -> None: 614 """ 615 Driver of the whole script, parses arguments, initialize the TUI and 616 the logger. 617 """ 618 parser = argparse.ArgumentParser(description='QMP TUI') 619 parser.add_argument('qmp_server', help='Address of the QMP server. ' 620 'Format <UNIX socket path | TCP addr:port>') 621 parser.add_argument('--num-retries', type=int, default=10, 622 help='Number of times to reconnect before giving up.') 623 parser.add_argument('--retry-delay', type=int, 624 help='Time(s) to wait before next retry. ' 625 'Default action is to wait 2s between each retry.') 626 parser.add_argument('--log-file', help='The Log file name') 627 parser.add_argument('--log-level', default='WARNING', 628 help='Log level <CRITICAL|ERROR|WARNING|INFO|DEBUG|>') 629 parser.add_argument('--asyncio-debug', action='store_true', 630 help='Enable debug mode for asyncio loop. ' 631 'Generates lot of output, makes TUI unusable when ' 632 'logs are logged in the TUI. ' 633 'Use only when logging to a file.') 634 args = parser.parse_args() 635 636 try: 637 address = QEMUMonitorProtocol.parse_address(args.qmp_server) 638 except QMPBadPortError as err: 639 parser.error(str(err)) 640 641 app = App(address, args.num_retries, args.retry_delay) 642 643 root_logger = logging.getLogger() 644 root_logger.setLevel(logging.getLevelName(args.log_level)) 645 646 if args.log_file: 647 root_logger.addHandler(logging.FileHandler(args.log_file)) 648 else: 649 root_logger.addHandler(TUILogHandler(app)) 650 651 app.run(args.asyncio_debug) 652 653 654if __name__ == '__main__': 655 main() 656