codekingpro/portable-devtools

1import asyncio2import contextlib3import mimetypes4import os.path5import shlex6import shutil7import stat8import subprocess9import sys10import tempfile11import threading12from typing import TypeVar13 14import urwid15from tornado.platform.asyncio import AddThreadSelectorEventLoop16 17from mitmproxy import addons18from mitmproxy import log19from mitmproxy import master20from mitmproxy import options21from mitmproxy.addons import errorcheck22from mitmproxy.addons import eventstore23from mitmproxy.addons import intercept24from mitmproxy.addons import readfile25from mitmproxy.addons import view26from mitmproxy.tools.console import consoleaddons27from mitmproxy.tools.console import defaultkeys28from mitmproxy.tools.console import keymap29from mitmproxy.tools.console import palettes30from mitmproxy.tools.console import signals31from mitmproxy.tools.console import window32from mitmproxy.utils import strutils33 34T = TypeVar("T", str, bytes)35 36 37class ConsoleMaster(master.Master):38    def __init__(self, opts: options.Options) -> None:39        super().__init__(opts)40 41        self.view: view.View = view.View()42        self.events = eventstore.EventStore()43        self.events.sig_add.connect(self.sig_add_log)44 45        self.stream_path = None46        self.keymap = keymap.Keymap(self)47        defaultkeys.map(self.keymap)48        self.options.errored.connect(self.options_error)49 50        self.addons.add(*addons.default_addons())51        self.addons.add(52            intercept.Intercept(),53            self.view,54            self.events,55            readfile.ReadFile(),56            consoleaddons.ConsoleAddon(self),57            keymap.KeymapConfig(self),58            errorcheck.ErrorCheck(repeat_errors_on_stderr=True),59        )60 61        self.window: window.Window | None = None62 63    def __setattr__(self, name, value):64        super().__setattr__(name, value)65        signals.update_settings.send()66 67    def options_error(self, exc) -> None:68        signals.status_message.send(message=str(exc), expire=1)69 70    def prompt_for_user_choice(self, prompt, callback) -> None:71        signals.status_prompt_onekey.send(72            prompt=prompt,73            keys=[74                ("yes", "y"),75                ("no", "n"),76            ],77            callback=callback,78        )79 80    def prompt_for_exit(self) -> None:81        self.prompt_for_user_choice("Quit", self.quit)82 83    def sig_add_log(self, entry: log.LogEntry):84        if log.log_tier(self.options.console_eventlog_verbosity) < log.log_tier(85            entry.level86        ):87            return88        if entry.level in ("error", "warn", "alert"):89            signals.status_message.send(90                message=(91                    entry.level,92                    f"{entry.level.title()}: {entry.msg.lstrip()}",93                ),94                expire=5,95            )96 97    def sig_call_in(self, seconds, callback):98        def cb(*_):99            return callback()100 101        self.loop.set_alarm_in(seconds, cb)102 103    @contextlib.contextmanager104    def uistopped(self):105        self.loop.stop()106        try:107            yield108        finally:109            self.loop.start()110            self.loop.screen_size = None111            self.loop.draw_screen()112 113    def get_editor(self) -> str:114        # based upon https://github.com/pallets/click/blob/main/src/click/_termui_impl.py115        if m := os.environ.get("MITMPROXY_EDITOR"):116            return m117        if m := os.environ.get("EDITOR"):118            return m119        for editor in "sensible-editor", "nano", "vim":120            if shutil.which(editor):121                return editor122        if os.name == "nt":123            return "notepad"124        else:125            return "vi"126 127    def get_hex_editor(self) -> str:128        editors = ["ghex", "hexedit", "hxd", "hexer", "hexcurse"]129        for editor in editors:130            if shutil.which(editor):131                return editor132        return self.get_editor()133 134    def spawn_editor(self, data: T) -> T:135        text = isinstance(data, str)136        fd, name = tempfile.mkstemp("", "mitmproxy", text=text)137        with_hexeditor = isinstance(data, bytes) and strutils.is_mostly_bin(data)138        with open(fd, "w" if text else "wb") as f:139            f.write(data)140        if with_hexeditor:141            c = self.get_hex_editor()142        else:143            c = self.get_editor()144        cmd = shlex.split(c)145        cmd.append(name)146        with self.uistopped():147            try:148                subprocess.call(cmd)149            except Exception:150                signals.status_message.send(message="Can't start editor: %s" % c)151            else:152                with open(name, "r" if text else "rb") as f:153                    data = f.read()154        os.unlink(name)155        return data156 157    def spawn_external_viewer(self, data, contenttype):158        if contenttype:159            contenttype = contenttype.split(";")[0]160            ext = mimetypes.guess_extension(contenttype) or ""161        else:162            ext = ""163        fd, name = tempfile.mkstemp(ext, "mproxy")164        os.write(fd, data)165        os.close(fd)166 167        # read-only to remind the user that this is a view function168        os.chmod(name, stat.S_IREAD)169 170        # hm which one should get priority?171        c = (172            os.environ.get("MITMPROXY_EDITOR")173            or os.environ.get("PAGER")174            or os.environ.get("EDITOR")175        )176        if not c:177            c = "less"178        cmd = shlex.split(c)179        cmd.append(name)180 181        with self.uistopped():182            try:183                subprocess.call(cmd, shell=False)184            except Exception:185                signals.status_message.send(186                    message="Can't start external viewer: %s" % " ".join(c)187                )188        # add a small delay before deletion so that the file is not removed before being loaded by the viewer189        t = threading.Timer(1.0, os.unlink, args=[name])190        t.start()191 192    def set_palette(self, *_) -> None:193        self.ui.register_palette(194            palettes.palettes[self.options.console_palette].palette(195                self.options.console_palette_transparent196            )197        )198        self.ui.clear()199 200    def inject_key(self, key):201        self.loop.process_input([key])202 203    async def running(self) -> None:204        if not sys.stdout.isatty():205            print(206                "Error: mitmproxy's console interface requires a tty. "207                "Please run mitmproxy in an interactive shell environment.",208                file=sys.stderr,209            )210            sys.exit(1)211 212        detected_encoding = urwid.detected_encoding.lower()213        if os.name != "nt" and detected_encoding and "utf" not in detected_encoding:214            print(215                f"mitmproxy expects a UTF-8 console environment, not {urwid.detected_encoding!r}. "216                f"Set your LANG environment variable to something like en_US.UTF-8.",217                file=sys.stderr,218            )219            # Experimental (04/2022): We just don't exit here and see if/how that affects users.220            # sys.exit(1)221        urwid.set_encoding("utf8")222 223        signals.call_in.connect(self.sig_call_in)224        self.ui = window.Screen()225        self.ui.set_terminal_properties(256)226        self.set_palette(None)227        self.options.subscribe(228            self.set_palette, ["console_palette", "console_palette_transparent"]229        )230 231        loop = asyncio.get_running_loop()232        if isinstance(loop, getattr(asyncio, "ProactorEventLoop", tuple())):233            # fix for https://bugs.python.org/issue37373234            loop = AddThreadSelectorEventLoop(loop)  # type: ignore235        self.loop = urwid.MainLoop(236            urwid.SolidFill("x"),237            event_loop=urwid.AsyncioEventLoop(loop=loop),238            screen=self.ui,239            handle_mouse=self.options.console_mouse,240        )241        self.window = window.Window(self)242        self.loop.widget = self.window243        self.window.refresh()244 245        self.loop.start()246 247        await super().running()248 249    async def done(self):250        self.loop.stop()251        await super().done()252 253    def overlay(self, widget, **kwargs):254        assert self.window255        self.window.set_overlay(widget, **kwargs)256 257    def switch_view(self, name):258        assert self.window259        self.window.push(name)260 261    def quit(self, a):262        if a != "n":263            self.shutdown()264