1import math
2
3import urwid
4
5from mitmproxy.tools.console import signals
6from mitmproxy.tools.console import grideditor
7from mitmproxy.tools.console import layoutwidget
8from mitmproxy.tools.console import keymap
9
10
11class SimpleOverlay(urwid.Overlay, layoutwidget.LayoutWidget):
12
13    def __init__(self, master, widget, parent, width, valign="middle"):
14        self.widget = widget
15        self.master = master
16        super().__init__(
17            widget,
18            parent,
19            align="center",
20            width=width,
21            valign=valign,
22            height="pack"
23        )
24
25    @property
26    def keyctx(self):
27        return getattr(self.widget, "keyctx")
28
29    def key_responder(self):
30        return self.widget.key_responder()
31
32    def focus_changed(self):
33        return self.widget.focus_changed()
34
35    def view_changed(self):
36        return self.widget.view_changed()
37
38    def layout_popping(self):
39        return self.widget.layout_popping()
40
41
42class Choice(urwid.WidgetWrap):
43    def __init__(self, txt, focus, current, shortcut):
44        if shortcut:
45            selection_type = "option_selected_key" if focus else "key"
46            txt = [(selection_type, shortcut), ") ", txt]
47        else:
48            txt = "   " + txt
49        if current:
50            s = "option_active_selected" if focus else "option_active"
51        else:
52            s = "option_selected" if focus else "text"
53        super().__init__(
54            urwid.AttrWrap(
55                urwid.Padding(urwid.Text(txt)),
56                s,
57            )
58        )
59
60    def selectable(self):
61        return True
62
63    def keypress(self, size, key):
64        return key
65
66
67class ChooserListWalker(urwid.ListWalker):
68    shortcuts = "123456789abcdefghijklmnoprstuvwxyz"
69
70    def __init__(self, choices, current):
71        self.index = 0
72        self.choices = choices
73        self.current = current
74
75    def _get(self, idx, focus):
76        c = self.choices[idx]
77        return Choice(c, focus, c == self.current, self.shortcuts[idx:idx + 1])
78
79    def set_focus(self, index):
80        self.index = index
81
82    def get_focus(self):
83        return self._get(self.index, True), self.index
84
85    def get_next(self, pos):
86        if pos >= len(self.choices) - 1:
87            return None, None
88        pos = pos + 1
89        return self._get(pos, False), pos
90
91    def get_prev(self, pos):
92        pos = pos - 1
93        if pos < 0:
94            return None, None
95        return self._get(pos, False), pos
96
97    def choice_by_shortcut(self, shortcut):
98        for i, choice in enumerate(self.choices):
99            if shortcut == self.shortcuts[i:i + 1]:
100                return choice
101        return None
102
103
104class Chooser(urwid.WidgetWrap, layoutwidget.LayoutWidget):
105    keyctx = "chooser"
106
107    def __init__(self, master, title, choices, current, callback):
108        self.master = master
109        self.choices = choices
110        self.callback = callback
111        choicewidth = max([len(i) for i in choices])
112        self.width = max(choicewidth, len(title)) + 7
113
114        self.walker = ChooserListWalker(choices, current)
115        super().__init__(
116            urwid.AttrWrap(
117                urwid.LineBox(
118                    urwid.BoxAdapter(
119                        urwid.ListBox(self.walker),
120                        len(choices)
121                    ),
122                    title=title
123                ),
124                "background"
125            )
126        )
127
128    def selectable(self):
129        return True
130
131    def keypress(self, size, key):
132        key = self.master.keymap.handle_only("chooser", key)
133        choice = self.walker.choice_by_shortcut(key)
134        if choice:
135            self.callback(choice)
136            signals.pop_view_state.send(self)
137            return
138        if key == "m_select":
139            self.callback(self.choices[self.walker.index])
140            signals.pop_view_state.send(self)
141            return
142        elif key in ["q", "esc"]:
143            signals.pop_view_state.send(self)
144            return
145
146        binding = self.master.keymap.get("global", key)
147        # This is extremely awkward. We need a better way to match nav keys only.
148        if binding and binding.command.startswith("console.nav"):
149            self.master.keymap.handle("global", key)
150        elif key in keymap.navkeys:
151            return super().keypress(size, key)
152
153
154class OptionsOverlay(urwid.WidgetWrap, layoutwidget.LayoutWidget):
155    keyctx = "grideditor"
156
157    def __init__(self, master, name, vals, vspace):
158        """
159            vspace: how much vertical space to keep clear
160        """
161        cols, rows = master.ui.get_cols_rows()
162        self.ge = grideditor.OptionsEditor(master, name, vals)
163        super().__init__(
164            urwid.AttrWrap(
165                urwid.LineBox(
166                    urwid.BoxAdapter(self.ge, rows - vspace),
167                    title=name
168                ),
169                "background"
170            )
171        )
172        self.width = math.ceil(cols * 0.8)
173
174    def key_responder(self):
175        return self.ge.key_responder()
176
177    def layout_popping(self):
178        return self.ge.layout_popping()
179
180
181class DataViewerOverlay(urwid.WidgetWrap, layoutwidget.LayoutWidget):
182    keyctx = "dataviewer"
183
184    def __init__(self, master, vals):
185        """
186            vspace: how much vertical space to keep clear
187        """
188        cols, rows = master.ui.get_cols_rows()
189        self.ge = grideditor.DataViewer(master, vals)
190        super().__init__(
191            urwid.AttrWrap(
192                urwid.LineBox(
193                    urwid.BoxAdapter(self.ge, rows - 5),
194                    title="Data viewer"
195                ),
196                "background"
197            )
198        )
199        self.width = math.ceil(cols * 0.8)
200
201    def key_responder(self):
202        return self.ge.key_responder()
203
204    def layout_popping(self):
205        return self.ge.layout_popping()
206