codekingpro/portable-devtools

1import asyncio2import importlib.machinery3import importlib.util4import logging5import os6import sys7import types8from collections.abc import Sequence9 10import mitmproxy.types as mtypes11from mitmproxy import addonmanager12from mitmproxy import command13from mitmproxy import ctx14from mitmproxy import eventsequence15from mitmproxy import exceptions16from mitmproxy import flow17from mitmproxy import hooks18from mitmproxy.utils import asyncio_utils19 20logger = logging.getLogger(__name__)21 22 23def load_script(path: str) -> types.ModuleType | None:24    fullname = "__mitmproxy_script__.{}".format(25        os.path.splitext(os.path.basename(path))[0]26    )27    # the fullname is not unique among scripts, so if there already is an existing script with said28    # fullname, remove it.29    sys.modules.pop(fullname, None)30    oldpath = sys.path31    sys.path.insert(0, os.path.dirname(path))32 33    try:34        loader = importlib.machinery.SourceFileLoader(fullname, path)35        spec = importlib.util.spec_from_loader(fullname, loader=loader)36        assert spec37        m = importlib.util.module_from_spec(spec)38        loader.exec_module(m)39        if not getattr(m, "name", None):40            m.name = path  # type: ignore41        return m42    except ImportError as e:43        if getattr(sys, "frozen", False):44            e.msg += (45                f".\n"46                f"Note that mitmproxy's binaries include their own Python environment. "47                f"If your addon requires the installation of additional dependencies, "48                f"please install mitmproxy from PyPI "49                f"(https://docs.mitmproxy.org/stable/overview-installation/#installation-from-the-python-package-index-pypi)."50            )51        script_error_handler(path, e)52        return None53    except Exception as e:54        script_error_handler(path, e)55        return None56    finally:57        sys.path[:] = oldpath58 59 60def script_error_handler(path: str, exc: Exception) -> None:61    """62    Log errors during script loading.63    """64    tback = exc.__traceback__65    tback = addonmanager.cut_traceback(66        tback, "invoke_addon_sync"67    )  # we're calling configure() on load68    tback = addonmanager.cut_traceback(69        tback, "_call_with_frames_removed"70    )  # module execution from importlib71    logger.error(f"error in script {path}", exc_info=(type(exc), exc, tback))72 73 74ReloadInterval = 175 76 77class Script:78    """79    An addon that manages a single script.80    """81 82    def __init__(self, path: str, reload: bool) -> None:83        self.name = "scriptmanager:" + path84        self.path = path85        self.fullpath = os.path.expanduser(path.strip("'\" "))86        self.ns: types.ModuleType | None = None87        self.is_running = False88 89        if not os.path.isfile(self.fullpath):90            raise exceptions.OptionsError(f"No such script: {self.fullpath}")91 92        self.reloadtask = None93        if reload:94            self.reloadtask = asyncio_utils.create_task(95                self.watcher(),96                name=f"script watcher for {path}",97                keep_ref=False,98            )99        else:100            self.loadscript()101 102    def running(self):103        self.is_running = True104 105    def done(self):106        if self.reloadtask:107            self.reloadtask.cancel()108 109    @property110    def addons(self):111        return [self.ns] if self.ns else []112 113    def loadscript(self):114        logger.info("Loading script %s" % self.path)115        if self.ns:116            ctx.master.addons.remove(self.ns)117        self.ns = None118        with addonmanager.safecall():119            ns = load_script(self.fullpath)120            ctx.master.addons.register(ns)121            self.ns = ns122        if self.ns:123            try:124                ctx.master.addons.invoke_addon_sync(125                    self.ns, hooks.ConfigureHook(ctx.options.keys())126                )127            except Exception as e:128                script_error_handler(self.fullpath, e)129            if self.is_running:130                # We're already running, so we call that on the addon now.131                ctx.master.addons.invoke_addon_sync(self.ns, hooks.RunningHook())132 133    async def watcher(self):134        # Script loading is terminally confused at the moment.135        # This here is a stopgap workaround to defer loading.136        await asyncio.sleep(0)137        last_mtime = 0.0138        while True:139            try:140                mtime = os.stat(self.fullpath).st_mtime141            except FileNotFoundError:142                logger.info("Removing script %s" % self.path)143                scripts = list(ctx.options.scripts)144                scripts.remove(self.path)145                ctx.options.update(scripts=scripts)146                return147            if mtime > last_mtime:148                self.loadscript()149                last_mtime = mtime150            await asyncio.sleep(ReloadInterval)151 152 153class ScriptLoader:154    """155    An addon that manages loading scripts from options.156    """157 158    def __init__(self):159        self.is_running = False160        self.addons = []161 162    def load(self, loader):163        loader.add_option("scripts", Sequence[str], [], "Execute a script.")164 165    def running(self):166        self.is_running = True167 168    @command.command("script.run")169    def script_run(self, flows: Sequence[flow.Flow], path: mtypes.Path) -> None:170        """171        Run a script on the specified flows. The script is configured with172        the current options and all lifecycle events for each flow are173        simulated. Note that the load event is not invoked.174        """175        if not os.path.isfile(path):176            logger.error("No such script: %s" % path)177            return178        mod = load_script(path)179        if mod:180            with addonmanager.safecall():181                ctx.master.addons.invoke_addon_sync(182                    mod,183                    hooks.ConfigureHook(ctx.options.keys()),184                )185                ctx.master.addons.invoke_addon_sync(mod, hooks.RunningHook())186                for f in flows:187                    for evt in eventsequence.iterate(f):188                        ctx.master.addons.invoke_addon_sync(mod, evt)189 190    def configure(self, updated):191        if "scripts" in updated:192            for s in ctx.options.scripts:193                if ctx.options.scripts.count(s) > 1:194                    raise exceptions.OptionsError("Duplicate script")195 196            for a in self.addons[:]:197                if a.path not in ctx.options.scripts:198                    logger.info("Un-loading script: %s" % a.path)199                    ctx.master.addons.remove(a)200                    self.addons.remove(a)201 202            # The machinations below are to ensure that:203            #   - Scripts remain in the same order204            #   - Scripts are not initialized un-necessarily. If only a205            #   script's order in the script list has changed, it is just206            #   moved.207 208            current = {}209            for a in self.addons:210                current[a.path] = a211 212            ordered = []213            newscripts = []214            for s in ctx.options.scripts:215                if s in current:216                    ordered.append(current[s])217                else:218                    sc = Script(s, True)219                    ordered.append(sc)220                    newscripts.append(sc)221 222            self.addons = ordered223 224            for s in newscripts:225                ctx.master.addons.register(s)226                if self.is_running:227                    # If we're already running, we configure and tell the addon228                    # we're up and running.229                    ctx.master.addons.invoke_addon_sync(s, hooks.RunningHook())230