codekingpro/portable-devtools

1import logging2import re3 4import urwid5 6from mitmproxy import flow7from mitmproxy.tools.console import commands8from mitmproxy.tools.console import common9from mitmproxy.tools.console import eventlog10from mitmproxy.tools.console import flowlist11from mitmproxy.tools.console import flowview12from mitmproxy.tools.console import grideditor13from mitmproxy.tools.console import help14from mitmproxy.tools.console import keybindings15from mitmproxy.tools.console import options16from mitmproxy.tools.console import overlay17from mitmproxy.tools.console import signals18from mitmproxy.tools.console import statusbar19 20 21class StackWidget(urwid.Frame):22    def __init__(self, window, widget, title, focus):23        self.is_focused = focus24        self.window = window25 26        if title:27            header = urwid.AttrMap(28                urwid.Text(title), "heading" if focus else "heading_inactive"29            )30        else:31            header = None32        super().__init__(widget, header=header)33 34    def mouse_event(self, size, event, button, col, row, focus):35        if event == "mouse press" and button == 1 and not self.is_focused:36            self.window.switch()37        return super().mouse_event(size, event, button, col, row, focus)38 39    def keypress(self, size, key):40        # Make sure that we don't propagate cursor events outside of the widget.41        # Otherwise, in a horizontal layout, urwid's Pile would change the focused widget42        # if we cannot scroll any further.43        ret = super().keypress(size, key)44        command = self._command_map[45            ret46        ]  # awkward as they don't implement a full dict api47        if command and command.startswith("cursor"):48            return None49        return ret50 51 52class WindowStack:53    def __init__(self, master, base):54        self.master = master55        self.windows = dict(56            flowlist=flowlist.FlowListBox(master),57            flowview=flowview.FlowView(master),58            commands=commands.Commands(master),59            keybindings=keybindings.KeyBindings(master),60            options=options.Options(master),61            help=help.HelpView(master),62            eventlog=eventlog.EventLog(master),63            edit_focus_query=grideditor.QueryEditor(master),64            edit_focus_comment=grideditor.CommentEditor(master),65            edit_focus_cookies=grideditor.CookieEditor(master),66            edit_focus_setcookies=grideditor.SetCookieEditor(master),67            edit_focus_setcookie_attrs=grideditor.CookieAttributeEditor(master),68            edit_focus_multipart_form=grideditor.RequestMultipartEditor(master),69            edit_focus_urlencoded_form=grideditor.RequestUrlEncodedEditor(master),70            edit_focus_path=grideditor.PathEditor(master),71            edit_focus_request_headers=grideditor.RequestHeaderEditor(master),72            edit_focus_response_headers=grideditor.ResponseHeaderEditor(master),73        )74        self.stack = [base]75        self.overlay = None76 77    def set_overlay(self, o, **kwargs):78        self.overlay = overlay.SimpleOverlay(79            self,80            o,81            self.top_widget(),82            o.width,83            **kwargs,84        )85 86    def top_window(self):87        """88        The current top window, ignoring overlays.89        """90        return self.windows[self.stack[-1]]91 92    def top_widget(self):93        """94        The current top widget - either a window or the active overlay.95        """96        if self.overlay:97            return self.overlay98        return self.top_window()99 100    def push(self, wname):101        if self.stack[-1] == wname:102            return103        prev = self.top_window()104        self.stack.append(wname)105        self.call("layout_pushed", prev)106 107    def pop(self, *args, **kwargs):108        """109        Pop off the stack, return True if we're already at the top.110        """111        if not self.overlay and len(self.stack) == 1:112            return True113        self.call("layout_popping")114        if self.overlay:115            self.overlay = None116        else:117            self.stack.pop()118 119    def call(self, name, *args, **kwargs):120        """121        Call a function on both the top window, and the overlay if there is122        one. If the widget has a key_responder, we call the function on the123        responder instead.124        """125        getattr(self.top_window(), name)(*args, **kwargs)126        if self.overlay:127            getattr(self.overlay, name)(*args, **kwargs)128 129 130class Window(urwid.Frame):131    def __init__(self, master):132        self.statusbar = statusbar.StatusBar(master)133        super().__init__(134            None, header=None, footer=urwid.AttrMap(self.statusbar, "background")135        )136        self.master = master137        self.master.view.sig_view_refresh.connect(self.view_changed)138        self.master.view.sig_view_add.connect(self.view_changed)139        self.master.view.sig_view_remove.connect(self.view_changed)140        self.master.view.sig_view_update.connect(self.view_changed)141        self.master.view.focus.sig_change.connect(self.view_changed)142        self.master.view.focus.sig_change.connect(self.focus_changed)143 144        signals.focus.connect(self.sig_focus)145        signals.flow_change.connect(self.flow_changed)146        signals.pop_view_state.connect(self.pop)147 148        self.master.options.subscribe(149            self.configure, ["console_layout", "console_layout_headers"]150        )151        self.pane = 0152        self.stacks = [WindowStack(master, "flowlist"), WindowStack(master, "eventlog")]153 154    def focus_stack(self):155        return self.stacks[self.pane]156 157    def configure(self, options, updated):158        self.refresh()159 160    def refresh(self):161        """162        Redraw the layout.163        """164        c = self.master.options.console_layout165        if c == "single":166            self.pane = 0167 168        def wrapped(idx):169            widget = self.stacks[idx].top_widget()170            if self.master.options.console_layout_headers:171                title = self.stacks[idx].top_window().title172            else:173                title = None174            return StackWidget(self, widget, title, self.pane == idx)175 176        w = None177        if c == "single":178            w = wrapped(0)179        elif c == "vertical":180            w = urwid.Pile(181                [wrapped(i) for i, s in enumerate(self.stacks)], focus_item=self.pane182            )183        else:184            w = urwid.Columns(185                [wrapped(i) for i, s in enumerate(self.stacks)],186                dividechars=1,187                focus_column=self.pane,188            )189 190        self.body = urwid.AttrMap(w, "background")191        signals.window_refresh.send()192 193    def flow_changed(self, flow: flow.Flow) -> None:194        if self.master.view.focus.flow:195            if flow.id == self.master.view.focus.flow.id:196                self.focus_changed()197 198    def focus_changed(self, *args, **kwargs):199        """200        Triggered when the focus changes - either when it's modified, or201        when it changes to a different flow altogether.202        """203        for i in self.stacks:204            i.call("focus_changed")205 206    def view_changed(self, *args, **kwargs):207        """208        Triggered when the view list has changed.209        """210        for i in self.stacks:211            i.call("view_changed")212 213    def set_overlay(self, o, **kwargs):214        """215        Set an overlay on the currently focused stack.216        """217        self.focus_stack().set_overlay(o, **kwargs)218        self.refresh()219 220    def push(self, wname):221        """222        Push a window onto the currently focused stack.223        """224        self.focus_stack().push(wname)225        self.refresh()226        self.view_changed()227        self.focus_changed()228 229    def pop(self) -> None:230        """231        Pop a window from the currently focused stack. If there is only one232        window on the stack, this prompts for exit.233        """234        if self.focus_stack().pop():235            self.master.prompt_for_exit()236        else:237            self.refresh()238            self.view_changed()239            self.focus_changed()240 241    def stacks_sorted_by_focus(self):242        """243        Returns:244            self.stacks, with the focused stack first.245        """246        stacks = self.stacks.copy()247        stacks.insert(0, stacks.pop(self.pane))248        return stacks249 250    def current(self, keyctx):251        """252        Returns the active widget with a matching key context, including overlays.253        If multiple stacks have an active widget with a matching key context,254        the currently focused stack is preferred.255        """256        for s in self.stacks_sorted_by_focus():257            t = s.top_widget()258            if t.keyctx == keyctx:259                return t260 261    def current_window(self, keyctx):262        """263        Returns the active window with a matching key context, ignoring overlays.264        If multiple stacks have an active widget with a matching key context,265        the currently focused stack is preferred.266        """267        for s in self.stacks_sorted_by_focus():268            t = s.top_window()269            if t.keyctx == keyctx:270                return t271 272    def sig_focus(self, section):273        self.focus_position = section274 275    def switch(self):276        """277        Switch between the two panes.278        """279        if self.master.options.console_layout == "single":280            self.pane = 0281        else:282            self.pane = (self.pane + 1) % len(self.stacks)283        self.refresh()284 285    def mouse_event(self, *args, **kwargs):286        # args: (size, event, button, col, row)287        k = super().mouse_event(*args, **kwargs)288        if not k:289            if args[1] == "mouse drag":290                signals.status_message.send(291                    message="Hold down fn, shift, alt or ctrl to select text or use the --set console_mouse=false parameter.",292                    expire=1,293                )294            elif args[1] == "mouse press" and args[2] == 4:295                self.keypress(args[0], "up")296            elif args[1] == "mouse press" and args[2] == 5:297                self.keypress(args[0], "down")298            else:299                return False300            return True301 302    def keypress(self, size, k):303        k = super().keypress(size, k)304        if k:305            return self.master.keymap.handle(self.focus_stack().top_widget().keyctx, k)306 307 308class Screen(urwid.raw_display.Screen):309    def __init__(self) -> None:310        super().__init__()311        self.logger = logging.getLogger("urwid")312 313    def write(self, data):314        if common.IS_WINDOWS_OR_WSL:315            # replace urwid's SI/SO, which produce artifacts under WSL.316            # at some point we may figure out what they actually do.317            data = re.sub("[\x0e\x0f]", "", data)318        super().write(data)319