text: import json
  import os
  from typing import Optional, Union, List, Tuple
  
  from PySide2.QtCore import Qt
  from PySide2.QtGui import QColor
  from PySide2.QtWidgets import QApplication, QFileDialog, QInputDialog, QLineEdit, QMessageBox
  
  import requests
  
  from ...config import Conf
  from ...data.object_container import ObjectContainer
  from ...logic import GlobalInfo
  from ...logic.threads import gui_thread_schedule_async
  from ...utils.io import isurl, download_url
  from ...errors import InvalidURLError, UnexpectedStatusCodeError
  from ..base_plugin import BasePlugin
  from .trace_statistics import TraceStatistics
  from .multi_trace import MultiTrace
  from .afl_qemu_bitmap import AFLQemuBitmap
  from .chess_trace_list import QChessTraceListDialog
  from .qtrace_viewer import QTraceViewer
  
  
  class TraceViewer(BasePlugin):
      def __init__(self, *args, **kwargs):
          super().__init__(*args, **kwargs)
  
          self.workspace.instance.register_container('trace', lambda: None, Optional[TraceStatistics],
                                                     'The current trace')
          self.workspace.instance.register_container('multi_trace', lambda: None, Optional[MultiTrace],
                                                     'The current set of multiple traces')
  
          # Register event callbacks
          # self.trace.am_subscribe(self._on_trace_updated)
          self.multi_trace.am_subscribe(self._on_trace_updated)
  
          self._viewers = []
  
      def teardown(self):
          # I don't really know a better way to do this. tbh allowing arbitrary widget additions is probably intractable
          for trace_viewer in self._viewers:
              trace_viewer.hide()
  
      #
      # Forwarding properties
      #
  
      @property
      def trace(self) -> Union[ObjectContainer, Optional[TraceStatistics]]:
          return self.workspace.instance.trace
  
      @property
      def multi_trace(self) -> Union[ObjectContainer, Optional[MultiTrace], AFLQemuBitmap]:
          return self.workspace.instance.multi_trace
  
      #
      # Event handlers
      #
  
      def _on_trace_updated(self):
          # redraw disassembly view
          view = self.workspace.view_manager.first_view_in_category('disassembly')
          if view is not None:
              view.redraw_current_graph()
          # refresh function table
          view = self.workspace.view_manager.first_view_in_category('functions')
          if view is not None:
              view.refresh()
  
      URL_ACTIONS = ['openbitmap', 'opentrace']
  
      def handle_url_action(self, action, kwargs):
          if action == 'openbitmap':
              try:
                  base = int(kwargs['base'], 16)
              except (ValueError, KeyError):
                  base = None
              gui_thread_schedule_async(GlobalInfo.main_window.bring_to_front)
              gui_thread_schedule_async(self.open_bitmap_multi_trace,
                                        args=(kwargs['path'],),
                                        kwargs={
                                            'base_addr': base,
                                        })
          elif action == 'opentrace':
              try:
                  base = int(kwargs['base'], 16)
              except (ValueError, KeyError):
                  base = None
              path = kwargs['path']
              gui_thread_schedule_async(GlobalInfo.main_window.bring_to_front)
              gui_thread_schedule_async(self.add_trace,
                                       kwargs={
                                           'trace_path': path,
                                           'base_addr': base,
                                       })
          else:
              raise ValueError("Trace plugin cannot handle url action " + action)
  
      #
      # features for the disassembly view
      #
  
      GRAPH_TRACE_LEGEND_WIDTH = 30
      GRAPH_TRACE_LEGEND_SPACING = 20
  
      def instrument_disassembly_view(self, dview):
          trace_viewer = QTraceViewer(self.workspace, dview, parent=dview)
          self._viewers.append(trace_viewer)
  
          dview.layout().addWidget(trace_viewer)
          trace_viewer.hide()
  
      def color_block(self, addr):
          if not self.multi_trace.am_none:
              if isinstance(self.multi_trace.am_obj, MultiTrace):
                  if not self.multi_trace.is_active_tab:
                      return None
              return self.multi_trace.get_hit_miss_color(addr)
          return None
  
      def handle_click_block(self, qblock, event):
          btn = event.button()
  
          if QApplication.keyboardModifiers() == Qt.ControlModifier and btn == Qt.RightButton:
              if self.multi_trace is not None and self.multi_trace.am_obj is not None:
                  the_trace = self.multi_trace.get_any_trace(qblock.addr)
                  if the_trace is not None:
                      self.trace.am_obj = TraceStatistics(self.workspace, the_trace, self.multi_trace.base_addr)
                      self.trace.am_event()
                      return True
  
          return False
  
      def draw_insn(self, qinsn, painter):
          # legend
          if not self.multi_trace.am_none:
              if isinstance(self.multi_trace.am_obj, MultiTrace):
                  if not self.multi_trace.is_active_tab:
                      return  # skip
          strata = self._gen_strata(qinsn.insn.addr)
          if strata is not None:
              legend_x = 0 - self.GRAPH_TRACE_LEGEND_WIDTH - self.GRAPH_TRACE_LEGEND_SPACING
              for (i, w) in strata:
                  color = self.trace.get_mark_color(qinsn.insn.addr, i)
                  painter.setPen(color)
                  painter.setBrush(color)
                  painter.drawRect(legend_x, 0, w, qinsn.height)
                  legend_x += w
  
      def _gen_strata(self, addr):
          if not self.trace.am_none:
              # count is cached in trace.
              count = self.trace.get_count(addr)
  
              if count > self.GRAPH_TRACE_LEGEND_WIDTH:
                  jump = count / self.GRAPH_TRACE_LEGEND_WIDTH
                  return [(int(jump * i), 1) for i in
                          range(self.GRAPH_TRACE_LEGEND_WIDTH)]
              elif count > 0:
                  width = self.GRAPH_TRACE_LEGEND_WIDTH // count
                  remainder = self.GRAPH_TRACE_LEGEND_WIDTH % count
                  return [(i, width + 1) for i in range(remainder)] + \
                         [(i, width) for i in range(remainder, count)]
  
          return None
  
      #
      # features for the functions view
      #
  
      def color_func(self, func):
          if not self.multi_trace.am_none:
              if isinstance(self.multi_trace.am_obj, MultiTrace):
                  if self.multi_trace.is_active_tab:
                      return self.multi_trace.get_percent_color(func)
              elif isinstance(self.multi_trace.am_obj, AFLQemuBitmap):
                  return self.multi_trace.get_percent_color(func)
              else:
                  # you should not reach here
                  raise RuntimeError("Impossible happened")
  
          if not self.trace.am_none:
              if func.addr in self.trace.func_addr_in_trace:
                  return QColor(0xf0, 0xe7, 0xda)
              return QColor(0xee, 0xee, 0xee)
          return None
  
      FUNC_COLUMNS = ('Coverage',)
  
      def extract_func_column(self, func, idx):
          assert idx == 0
          if self.multi_trace.am_none:
              cov = 0
              rend = ''
          else:
              cov = self.multi_trace.get_coverage(func)
              rend = "%.2f%%" % cov
  
          return cov, rend
  
      #
      # features for loading traces!
      #
  
      MENU_BUTTONS = [
          'Open/Add trace...',
          'Clear trace',
          'Open AFL bitmap...',
          'Open inverted AFL bitmap...',
          'Reset AFL bitmap',
          'Open traces from CHECRS...',
      ]
      ADD_TRACE_ID = 0
      RESET_TRACE_ID = 1
      OPEN_AFL_BITMAP_ID = 2
      OPEN_AFL_BITMAP_INVERTED_ID = 3
      RESET_AFL_BITMAP = 4
      OPEN_TRACES_FROM_CHECRS = 5
  
      def handle_click_menu(self, idx):
  
          if idx < 0 or idx >= len(self.MENU_BUTTONS):
              return
  
          if self.workspace.instance.project.am_none:
              return
  
          mapping = {
              self.ADD_TRACE_ID: self.add_trace,
              self.RESET_TRACE_ID: self.reset_trace,
              self.OPEN_AFL_BITMAP_ID: self.open_bitmap_multi_trace,
              self.OPEN_AFL_BITMAP_INVERTED_ID: self.open_inverted_bitmap_multi_trace,
              self.RESET_AFL_BITMAP: self.reset_bitmap,
              self.OPEN_TRACES_FROM_CHECRS: self.open_traces_from_checrs,
          }
  
          mapping.get(idx)()
  
      # def open_trace(self):
      #     trace, base_addr = self._open_json_trace_dialog()
      #     if trace is None or base_addr is None:
      #         return
  
      #     # self.trace.am_obj = TraceStatistics(self.workspace, trace, base_addr)
      #     # self.trace.am_event()
      #     self.multi_trace.am_obj = MultiTrace(self.workspace)
      #     self.trace.am_obj = self.multi_trace.am_obj.add_trace(trace, base_addr)
      #     self.multi_trace.am_event()
      #     self.trace.am_event()
  
      def add_trace(self, trace_path=None, base_addr=None):
          trace, base_addr = self._open_trace(trace_path, base_addr)
          if trace is None or base_addr is None:
              return
  
          self._add_trace(trace, base_addr)
  
      def reset_trace(self):
          self.trace.am_obj = None
          self.trace.am_event()
          self.multi_trace.am_obj = None
          self.multi_trace.am_event()
  
      def open_bitmap_multi_trace(self, trace_path=None, base_addr=None):
          r = self._open_bitmap_multi_trace(trace_path, base_addr)
          if r is None:
              return
          trace, base_addr = r
          self.multi_trace.am_obj = AFLQemuBitmap(self.workspace, trace, base_addr)
          self.multi_trace.am_event()
  
      def open_inverted_bitmap_multi_trace(self, trace_path=None, base_addr=None):
          r = self._open_bitmap_multi_trace(trace_path, base_addr)
          if r is None:
              return
          trace, base_addr = r
          self.multi_trace.am_obj = AFLQemuBitmap(self.workspace, trace, base_addr,
                                                  bits_inverted=True)
          self.multi_trace.am_event()
  
      def reset_bitmap(self):
          self.multi_trace.am_obj = None
          self.multi_trace.am_event()
  
      def _open_bitmap_multi_trace(self, trace_path, base_addr):
  
          if trace_path is None:
              trace_path = self._open_trace_dialog(tfilter='')
              if trace_path is None:
                  return None
  
          if base_addr is None:
              base_addr = self._open_baseaddr_dialog(0x4000000000)
              if base_addr is None:
                  return None
  
          if isurl(trace_path):
              try:
                  trace = download_url(trace_path, parent=self.workspace._main_window, to_file=False)
              except InvalidURLError:
                  QMessageBox.critical(self.workspace._main_window,
                                       "Downloading failed",
                                       "angr management failed to download the file. The URL is invalid.")
                  return None
              except UnexpectedStatusCodeError as ex:
                  QMessageBox.critical(self.workspace._main_window,
                                       "Downloading failed",
                                       "angr management failed to retrieve the header of the file. "
                                       "The HTTP request returned an unexpected status code %d." % ex.status_code)
                  return None
          else:
              with open(trace_path, 'rb') as f:
                  trace = f.read()
  
          return trace, base_addr
  
      def _open_trace(self, trace_path, base_addr):
          if trace_path is None:
              return self._open_json_trace_dialog()
          else:
              if isurl(trace_path):
                  try:
                      trace_bytes = download_url(trace_path, parent=self.workspace._main_window,
                                           to_file=False)
                      trace = json.loads(trace_bytes)
                  except InvalidURLError:
                      QMessageBox.critical(self.workspace._main_window,
                                           "Downloading failed",
                                           "angr management failed to download the file. The URL is invalid.")
                      trace = None
                  except UnexpectedStatusCodeError as ex:
                      QMessageBox.critical(self.workspace._main_window,
                                           "Downloading failed",
                                           "angr management failed to retrieve the header of the file. "
                                           "The HTTP request returned an unexpected status code %d." % ex.status_code)
                      trace = None
              else:
                  with open(trace_path, 'r') as f:
                      trace = json.load(f)
  
              if base_addr is None:
                  base_addr = self._open_baseaddr_dialog(0x4000000000)
  
              return trace, base_addr
  
      def _open_trace_dialog(self, tfilter):
          file_path, _ = QFileDialog.getOpenFileName(None, "Open a trace", "", tfilter)
          try:
              with open(file_path, 'rb') as f:
                  f.read(1)
          except FileNotFoundError:
              return None
  
          return file_path
  
      def _open_baseaddr_dialog(self, default_base):
          base_addr, _ = QInputDialog.getText(None, "Input Trace Base Address",
                                              "Base Address:", QLineEdit.Normal,
                                              hex(default_base))
          try:
              return int(base_addr, 16)
          except ValueError:
              return None
  
      def _open_json_trace_dialog(self) -> Tuple[Optional[List[int]], Optional[int]]:
          # project = self.workspace.instance.project
          trace_file_name = self._open_trace_dialog(tfilter='json (*.json)')
  
          if trace_file_name is None:
              return None, None
  
          with open(trace_file_name, 'r') as f:
              trace = json.load(f)
  
          if not isinstance(trace, dict):
              QMessageBox.critical(self.workspace._main_window,
                                   "Incorrect trace format",
                                   "Failed to open the JSON trace. We expect the JSON trace to be a dict.")
              return None, None
          elif "bb_addrs" not in trace.keys():
              QMessageBox.critical(self.workspace._main_window,
                                   "Incorrect trace format",
                                   "Failed to open the JSON trace."
                                   " We expect the JSON trace to contain the field 'bb_addrs'.")
              return None, None
          elif not isinstance(trace["bb_addrs"], list):
              QMessageBox.critical(self.workspace._main_window,
                                   "Incorrect trace format",
                                   "Failed to open the JSON trace."
                                   "We expect the JSON trace bb_addrs field to be a list of integers.")
              return None, None
  
          base_addr = self._open_baseaddr_dialog(0x4000000000)
  
          return trace, base_addr
  
      def _add_trace(self, trace, base_addr):
          if self.multi_trace.am_obj is None:
              self.multi_trace.am_obj = MultiTrace(self.workspace)
          self.trace.am_obj = self.multi_trace.am_obj.add_trace(trace, base_addr)
          self.multi_trace.am_event()
          self.trace.am_event()
  
      def open_traces_from_checrs(self):
  
          if not Conf.checrs_rest_endpoint_url:
              QMessageBox.critical(self.workspace.main_window,
                                   "Unspecified CHECRS REST endpoint",
                                   "The CHECRS REST endpoint is not specified.")
              return
  
          connector = self.workspace.plugins.get_plugin_instance_by_name("ChessConnector")
          if connector is None:
              return
          target_image_id = connector.target_image_id
          if not target_image_id:
              return
  
          dialog = QChessTraceListDialog(self.workspace, parent=self.workspace.main_window)
          dialog.exec_()
  
          traces = [ ]
          failures: List[Tuple[str,str]] = [ ]
  
          if dialog.trace_ids:
              for input_id, trace_id in dialog.trace_ids:
                  # talk to the backend to get the trace
                  try:
                      url = os.path.join(Conf.checrs_rest_endpoint_url,
                                         f"v1/targets/{target_image_id}/seeds/{input_id}/trace")
                      r = requests.get(url)
                  except Exception as ex:
                      failures.append((trace_id, str(ex)))
                      continue
                  data = r.text
                  try:
                      trace = json.loads(data)
                  except json.JSONDecodeError as ex:
                      failures.append((trace_id, str(ex)))
                      continue
                  traces.append(trace)
  
          for trace in traces:
              self._add_trace(trace, 0x4000000000)
  
          if failures:
              msg = [ ]
              for trace_id, exc in failures:
                  msg.append(f"Trace {trace_id}: {exc}.")
              QMessageBox.warning(self.workspace.main_window,
                                  "Failed to load some traces",
                                  "angr management failed to load some traces.\n"
                                  "Details:\n"
                                  + "\n".join(msg))
  