Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes15kdownloads
consoleaddons.py764 linesDownload Raw Back to console
1import csv2import logging3import re4from collections.abc import Sequence5 6import mitmproxy.types7from mitmproxy import command8from mitmproxy import command_lexer9from mitmproxy import contentviews10from mitmproxy import ctx11from mitmproxy import dns12from mitmproxy import exceptions13from mitmproxy import flow14from mitmproxy import http15from mitmproxy import log16from mitmproxy import tcp17from mitmproxy import udp18from mitmproxy.contentviews import ContentviewMessage19from mitmproxy.exceptions import CommandError20from mitmproxy.log import ALERT21from mitmproxy.tools.console import keymap22from mitmproxy.tools.console import overlay23from mitmproxy.tools.console import signals24from mitmproxy.utils import strutils25 26logger = logging.getLogger(__name__)27 28 29console_palettes = [30    "lowlight",31    "lowdark",32    "light",33    "dark",34    "solarized_light",35    "solarized_dark",36]37console_layouts = [38    "single",39    "vertical",40    "horizontal",41]42 43console_flowlist_layout = ["default", "table", "list"]44 45 46class ConsoleAddon:47    """48    An addon that exposes console-specific commands, and hooks into required49    events.50    """51 52    def __init__(self, master):53        self.master = master54        self.started = False55 56    def load(self, loader):57        loader.add_option(58            "console_default_contentview",59            str,60            "auto",61            "The default content view mode.",62            choices=contentviews.registry.available_views(),63        )64        loader.add_option(65            "console_eventlog_verbosity",66            str,67            "info",68            "EventLog verbosity.",69            choices=log.LogLevels,70        )71        loader.add_option(72            "console_layout",73            str,74            "single",75            "Console layout.",76            choices=sorted(console_layouts),77        )78        loader.add_option(79            "console_layout_headers",80            bool,81            True,82            "Show layout component headers",83        )84        loader.add_option(85            "console_focus_follow", bool, False, "Focus follows new flows."86        )87        loader.add_option(88            "console_palette",89            str,90            "solarized_dark",91            "Color palette.",92            choices=sorted(console_palettes),93        )94        loader.add_option(95            "console_palette_transparent",96            bool,97            True,98            "Set transparent background for palette.",99        )100        loader.add_option("console_mouse", bool, True, "Console mouse interaction.")101        loader.add_option(102            "console_flowlist_layout",103            str,104            "default",105            "Set the flowlist layout",106            choices=sorted(console_flowlist_layout),107        )108        loader.add_option(109            "console_strip_trailing_newlines",110            bool,111            False,112            "Strip trailing newlines from edited request/response bodies.",113        )114        loader.add_option(115            "console_quickhelp_visible",116            bool,117            True,118            "Show the quick help bar at the bottom of the console.",119        )120 121    @command.command("console.layout.options")122    def layout_options(self) -> Sequence[str]:123        """124        Returns the available options for the console_layout option.125        """126        return ["single", "vertical", "horizontal"]127 128    @command.command("console.layout.cycle")129    def layout_cycle(self) -> None:130        """131        Cycle through the console layout options.132        """133        opts = self.layout_options()134        off = self.layout_options().index(ctx.options.console_layout)135        ctx.options.update(console_layout=opts[(off + 1) % len(opts)])136 137    @command.command("console.panes.next")138    def panes_next(self) -> None:139        """140        Go to the next layout pane.141        """142        self.master.window.switch()143 144    @command.command("console.panes.prev")145    def panes_prev(self) -> None:146        """147        Go to the previous layout pane.148        """149        return self.panes_next()150 151    @command.command("console.options.reset.focus")152    def options_reset_current(self) -> None:153        """154        Reset the current option in the options editor.155        """156        fv = self.master.window.current("options")157        if not fv:158            raise exceptions.CommandError("Not viewing options.")159        self.master.commands.call_strings("options.reset.one", [fv.current_name()])160 161    @command.command("console.nav.start")162    def nav_start(self) -> None:163        """164        Go to the start of a list or scrollable.165        """166        self.master.inject_key("m_start")167 168    @command.command("console.nav.end")169    def nav_end(self) -> None:170        """171        Go to the end of a list or scrollable.172        """173        self.master.inject_key("m_end")174 175    @command.command("console.nav.next")176    def nav_next(self) -> None:177        """178        Go to the next navigatable item.179        """180        self.master.inject_key("m_next")181 182    @command.command("console.nav.select")183    def nav_select(self) -> None:184        """185        Select a navigable item for viewing or editing.186        """187        self.master.inject_key("m_select")188 189    @command.command("console.nav.up")190    def nav_up(self) -> None:191        """192        Go up.193        """194        self.master.inject_key("up")195 196    @command.command("console.nav.down")197    def nav_down(self) -> None:198        """199        Go down.200        """201        self.master.inject_key("down")202 203    @command.command("console.nav.pageup")204    def nav_pageup(self) -> None:205        """206        Go up.207        """208        self.master.inject_key("page up")209 210    @command.command("console.nav.pagedown")211    def nav_pagedown(self) -> None:212        """213        Go down.214        """215        self.master.inject_key("page down")216 217    @command.command("console.nav.left")218    def nav_left(self) -> None:219        """220        Go left.221        """222        self.master.inject_key("left")223 224    @command.command("console.nav.right")225    def nav_right(self) -> None:226        """227        Go right.228        """229        self.master.inject_key("right")230 231    @command.command("console.choose")232    def console_choose(233        self,234        prompt: str,235        choices: Sequence[str],236        cmd: mitmproxy.types.Cmd,237        *args: mitmproxy.types.CmdArgs,238    ) -> None:239        """240        Prompt the user to choose from a specified list of strings, then241        invoke another command with all occurrences of {choice} replaced by242        the choice the user made.243        """244 245        def callback(opt):246            # We're now outside of the call context...247            repl = [arg.replace("{choice}", opt) for arg in args]248            try:249                self.master.commands.call_strings(cmd, repl)250            except exceptions.CommandError as e:251                logger.error(str(e))252 253        self.master.overlay(overlay.Chooser(self.master, prompt, choices, "", callback))254 255    @command.command("console.choose.cmd")256    def console_choose_cmd(257        self,258        prompt: str,259        choicecmd: mitmproxy.types.Cmd,260        subcmd: mitmproxy.types.Cmd,261        *args: mitmproxy.types.CmdArgs,262    ) -> None:263        """264        Prompt the user to choose from a list of strings returned by a265        command, then invoke another command with all occurrences of {choice}266        replaced by the choice the user made.267        """268        choices = ctx.master.commands.execute(choicecmd)269 270        def callback(opt):271            # We're now outside of the call context...272            repl = [arg.replace("{choice}", opt) for arg in args]273            try:274                self.master.commands.call_strings(subcmd, repl)275            except exceptions.CommandError as e:276                logger.error(str(e))277 278        self.master.overlay(overlay.Chooser(self.master, prompt, choices, "", callback))279 280    @command.command("console.command")281    def console_command(self, *command_str: str) -> None:282        """283        Prompt the user to edit a command with a (possibly empty) starting value.284        """285        quoted = " ".join(command_lexer.quote(x) for x in command_str)286        if quoted:287            quoted += " "288        signals.status_prompt_command.send(partial=quoted)289 290    @command.command("console.command.confirm")291    def console_command_confirm(292        self,293        prompt: str,294        cmd: mitmproxy.types.Cmd,295        *args: mitmproxy.types.CmdArgs,296    ) -> None:297        """298        Prompt the user before running the specified command.299        """300 301        def callback(opt):302            if opt == "n":303                return304            try:305                self.master.commands.call_strings(cmd, args)306            except exceptions.CommandError as e:307                logger.exception(str(e))308 309        self.master.prompt_for_user_choice(prompt, callback)310 311    @command.command("console.command.set")312    def console_command_set(self, option_name: str) -> None:313        """314        Prompt the user to set an option.315        """316        option_value = getattr(self.master.options, option_name, None) or ""317        set_command = f"set {option_name} {option_value!r}"318        cursor = len(set_command) - 1319        signals.status_prompt_command.send(partial=set_command, cursor=cursor)320 321    @command.command("console.view.keybindings")322    def view_keybindings(self) -> None:323        """View the commands list."""324        self.master.switch_view("keybindings")325 326    @command.command("console.view.commands")327    def view_commands(self) -> None:328        """View the commands list."""329        self.master.switch_view("commands")330 331    @command.command("console.view.options")332    def view_options(self) -> None:333        """View the options editor."""334        self.master.switch_view("options")335 336    @command.command("console.view.eventlog")337    def view_eventlog(self) -> None:338        """View the event log."""339        self.master.switch_view("eventlog")340 341    @command.command("console.view.help")342    def view_help(self) -> None:343        """View help."""344        self.master.switch_view("help")345 346    @command.command("console.view.flow")347    def view_flow(self, flow: flow.Flow) -> None:348        """View a flow."""349        if isinstance(flow, (http.HTTPFlow, tcp.TCPFlow, udp.UDPFlow, dns.DNSFlow)):350            self.master.switch_view("flowview")351        else:352            logger.warning(f"No detail view for {type(flow).__name__}.")353 354    @command.command("console.exit")355    def exit(self) -> None:356        """Exit mitmproxy."""357        self.master.shutdown()358 359    @command.command("console.view.pop")360    def view_pop(self) -> None:361        """362        Pop a view off the console stack. At the top level, this prompts the363        user to exit mitmproxy.364        """365        signals.pop_view_state.send()366 367    @command.command("console.bodyview")368    @command.argument("part", type=mitmproxy.types.Choice("console.bodyview.options"))369    def bodyview(self, flow: flow.Flow, part: str) -> None:370        """371        Spawn an external viewer for a flow request or response body based372        on the detected MIME type. We use the mailcap system to find the373        correct viewer, and fall back to the programs in $PAGER or $EDITOR374        if necessary.375        """376        fpart = getattr(flow, part, None)377        if not fpart:378            raise exceptions.CommandError(379                "Part must be either request or response, not %s." % part380            )381        t = fpart.headers.get("content-type")382        content = fpart.get_content(strict=False)383        if not content:384            raise exceptions.CommandError("No content to view.")385        self.master.spawn_external_viewer(content, t)386 387    @command.command("console.bodyview.options")388    def bodyview_options(self) -> Sequence[str]:389        """390        Possible parts for console.bodyview.391        """392        return ["request", "response"]393 394    @command.command("console.edit.focus.options")395    def edit_focus_options(self) -> Sequence[str]:396        """397        Possible components for console.edit.focus.398        """399        flow = self.master.view.focus.flow400        focus_options = []401 402        try:403            view_name = self.master.commands.call("console.flowview.mode")404        except CommandError:405            view_name = "auto"406 407        def add_message_edit_option(408            message_name: str, message: ContentviewMessage | None409        ) -> None:410            if message is None:411                return412            data, _ = contentviews.get_data(message)413            cv = contentviews.registry.get_view(414                data or b"",415                contentviews.make_metadata(message, flow),416                view_name,417            )418            if isinstance(cv, contentviews.InteractiveContentview):419                focus_options.append(f"{message_name} ({cv.name})")420 421        if flow is None:422            raise exceptions.CommandError("No flow selected.")423 424        if isinstance(flow, tcp.TCPFlow):425            focus_options.append("tcp-message")426            add_message_edit_option("tcp-message", flow.messages[-1])427        elif isinstance(flow, udp.UDPFlow):428            focus_options.append("udp-message")429            add_message_edit_option("udp-message", flow.messages[-1])430        elif isinstance(flow, http.HTTPFlow):431            focus_options.extend(432                [433                    "cookies",434                    "urlencoded form",435                    "multipart form",436                    "path",437                    "method",438                    "query",439                    "reason",440                    "request-headers",441                    "response-headers",442                    "request-body",443                    "response-body",444                    "status_code",445                    "set-cookies",446                    "url",447                ]448            )449            add_message_edit_option("request-body", flow.request)450            add_message_edit_option("response-body", flow.response)451            if flow.websocket:452                add_message_edit_option(453                    "websocket-message", flow.websocket.messages[-1]454                )455        elif isinstance(flow, dns.DNSFlow):456            raise exceptions.CommandError(457                "Cannot edit DNS flows yet, please submit a patch."458            )459 460        focus_options.append("comment")461 462        return focus_options463 464    @command.command("console.edit.focus")465    @command.argument(466        "flow_part", type=mitmproxy.types.Choice("console.edit.focus.options")467    )468    def edit_focus(self, flow_part: str) -> None:469        """470        Edit a component of the currently focused flow.471        """472        flow = self.master.view.focus.flow473        # This shouldn't be necessary once this command is "console.edit @focus",474        # but for now it is.475        if not flow:476            raise exceptions.CommandError("No flow selected.")477        flow.backup()478 479        require_dummy_response = (480            flow_part in ("response-headers", "response-body", "set-cookies")481            and flow.response is None482        )483        if require_dummy_response:484            flow.response = http.Response.make()485        if flow_part == "comment":486            self.master.switch_view("edit_focus_comment")487        elif flow_part == "cookies":488            self.master.switch_view("edit_focus_cookies")489        elif flow_part == "urlencoded form":490            self.master.switch_view("edit_focus_urlencoded_form")491        elif flow_part == "multipart form":492            self.master.switch_view("edit_focus_multipart_form")493        elif flow_part == "path":494            self.master.switch_view("edit_focus_path")495        elif flow_part == "query":496            self.master.switch_view("edit_focus_query")497        elif flow_part == "request-headers":498            self.master.switch_view("edit_focus_request_headers")499        elif flow_part == "response-headers":500            self.master.switch_view("edit_focus_response_headers")501        elif m := re.match(502            r"(?P<message>(request|response)-body|(tcp|udp|websocket)-message) \((?P<contentview>.+)\)",503            flow_part,504        ):505            match m["message"]:506                case "request-body":507                    message = flow.request508                case "response-body":509                    message = flow.response510                case "tcp-message" | "udp-message":511                    message = flow.messages[-1]512                case "websocket-message":513                    message = flow.websocket.messages[-1]514                case _:515                    assert False, "should be exhaustive"516 517            cv = contentviews.registry.get(m["contentview"])518            if not cv or not isinstance(cv, contentviews.InteractiveContentview):519                raise CommandError(520                    f"Contentview {m['contentview']} is not bidirectional."521                )522 523            pretty = contentviews.prettify_message(message, flow, cv.name)524            prettified = self.master.spawn_editor(pretty.text)525 526            message.content = contentviews.reencode_message(527                prettified,528                message,529                flow,530                cv.name,531            )532 533        elif flow_part in ("request-body", "response-body"):534            if flow_part == "request-body":535                message = flow.request536            else:537                message = flow.response538            c = self.master.spawn_editor(message.get_content(strict=False) or b"")539            # Many editors make it hard to save a file without a terminating540            # newline on the last line. When editing message bodies, this can541            # cause problems. We strip trailing newlines by default, but this542            # behavior is configurable.543            if self.master.options.console_strip_trailing_newlines:544                c = c.rstrip(b"\n")545            message.content = c546        elif flow_part == "set-cookies":547            self.master.switch_view("edit_focus_setcookies")548        elif flow_part == "url":549            url = flow.request.url.encode()550            edited_url = self.master.spawn_editor(url)551            url = edited_url.rstrip(b"\n")552            flow.request.url = url.decode()553        elif flow_part in ["method", "status_code", "reason"]:554            self.master.commands.call_strings(555                "console.command", ["flow.set", "@focus", flow_part]556            )557        elif flow_part in ["tcp-message", "udp-message", "websocket-message"]:558            if flow_part == "websocket-message":559                message = flow.websocket.messages[-1]560            else:561                message = flow.messages[-1]562            c = self.master.spawn_editor(message.content or b"")563            if self.master.options.console_strip_trailing_newlines:564                c = c.rstrip(b"\n")565            message.content = c.rstrip(b"\n")566 567    def _grideditor(self):568        gewidget = self.master.window.current("grideditor")569        if not gewidget:570            raise exceptions.CommandError("Not in a grideditor.")571        return gewidget.key_responder()572 573    @command.command("console.grideditor.add")574    def grideditor_add(self) -> None:575        """576        Add a row after the cursor.577        """578        self._grideditor().cmd_add()579 580    @command.command("console.grideditor.insert")581    def grideditor_insert(self) -> None:582        """583        Insert a row before the cursor.584        """585        self._grideditor().cmd_insert()586 587    @command.command("console.grideditor.delete")588    def grideditor_delete(self) -> None:589        """590        Delete row591        """592        self._grideditor().cmd_delete()593 594    @command.command("console.grideditor.load")595    def grideditor_load(self, path: mitmproxy.types.Path) -> None:596        """597        Read a file into the current cell.598        """599        self._grideditor().cmd_read_file(path)600 601    @command.command("console.grideditor.load_escaped")602    def grideditor_load_escaped(self, path: mitmproxy.types.Path) -> None:603        """604        Read a file containing a Python-style escaped string into the605        current cell.606        """607        self._grideditor().cmd_read_file_escaped(path)608 609    @command.command("console.grideditor.save")610    def grideditor_save(self, path: mitmproxy.types.Path) -> None:611        """612        Save data to file as a CSV.613        """614        rows = self._grideditor().value615        try:616            with open(path, "w", newline="", encoding="utf8") as fp:617                writer = csv.writer(fp)618                for row in rows:619                    writer.writerow(620                        [strutils.always_str(x) or "" for x in row]  # type: ignore621                    )622            logger.log(ALERT, "Saved %s rows as CSV." % (len(rows)))623        except OSError as e:624            logger.error(str(e))625 626    @command.command("console.grideditor.editor")627    def grideditor_editor(self) -> None:628        """629        Spawn an external editor on the current cell.630        """631        self._grideditor().cmd_spawn_editor()632 633    @command.command("console.flowview.mode.set")634    @command.argument(635        "mode", type=mitmproxy.types.Choice("console.flowview.mode.options")636    )637    def flowview_mode_set(self, mode: str) -> None:638        """639        Set the display mode for the current flow view.640        """641        fv = self.master.window.current_window("flowview")642        if not fv:643            raise exceptions.CommandError("Not viewing a flow.")644        idx = fv.body.tab_offset645 646        if mode.lower() not in self.flowview_mode_options():647            raise exceptions.CommandError("Invalid flowview mode.")648 649        try:650            self.master.commands.call_strings(651                "view.settings.setval", ["@focus", f"flowview_mode_{idx}", mode]652            )653        except exceptions.CommandError as e:654            logger.error(str(e))655 656    @command.command("console.flowview.mode.options")657    def flowview_mode_options(self) -> Sequence[str]:658        """659        Returns the valid options for the flowview mode.660        """661        return contentviews.registry.available_views()662 663    @command.command("console.flowview.mode")664    def flowview_mode(self) -> str:665        """666        Get the display mode for the current flow view.667        """668        fv = self.master.window.current_window("flowview")669        if not fv:670            raise exceptions.CommandError("Not viewing a flow.")671        idx = fv.body.tab_offset672 673        return self.master.commands.call_strings(674            "view.settings.getval",675            [676                "@focus",677                f"flowview_mode_{idx}",678                self.master.options.console_default_contentview,679            ],680        )681 682    @command.command("console.key.contexts")683    def key_contexts(self) -> Sequence[str]:684        """685        The available contexts for key binding.686        """687        return list(sorted(keymap.Contexts))688 689    @command.command("console.key.bind")690    def key_bind(691        self,692        contexts: Sequence[str],693        key: str,694        cmd: mitmproxy.types.Cmd,695        *args: mitmproxy.types.CmdArgs,696    ) -> None:697        """698        Bind a shortcut key.699        """700        try:701            self.master.keymap.add(key, cmd + " " + " ".join(args), contexts, "")702        except ValueError as v:703            raise exceptions.CommandError(v)704 705    @command.command("console.key.unbind")706    def key_unbind(self, contexts: Sequence[str], key: str) -> None:707        """708        Un-bind a shortcut key.709        """710        try:711            self.master.keymap.remove(key, contexts)712        except ValueError as v:713            raise exceptions.CommandError(v)714 715    def _keyfocus(self):716        kwidget = self.master.window.current("keybindings")717        if not kwidget:718            raise exceptions.CommandError("Not viewing key bindings.")719        f = kwidget.get_focused_binding()720        if not f:721            raise exceptions.CommandError("No key binding focused")722        return f723 724    @command.command("console.key.unbind.focus")725    def key_unbind_focus(self) -> None:726        """727        Un-bind the shortcut key currently focused in the key binding viewer.728        """729        b = self._keyfocus()730        try:731            self.master.keymap.remove(b.key, b.contexts)732        except ValueError as v:733            raise exceptions.CommandError(v)734 735    @command.command("console.key.execute.focus")736    def key_execute_focus(self) -> None:737        """738        Execute the currently focused key binding.739        """740        b = self._keyfocus()741        self.console_command(b.command)742 743    @command.command("console.key.edit.focus")744    def key_edit_focus(self) -> None:745        """746        Execute the currently focused key binding.747        """748        b = self._keyfocus()749        self.console_command(750            "console.key.bind",751            ",".join(b.contexts),752            b.key,753            b.command,754        )755 756    def running(self):757        self.started = True758 759    def update(self, flows) -> None:760        if not flows:761            signals.update_settings.send()762        for f in flows:763            signals.flow_change.send(flow=f)764 
codekingpro/portable-devtools · Team Ai