codekingpro/portable-devtools

1"""2The View:3 4- Keeps track of a store of flows5- Maintains a filtered, ordered view onto that list of flows6- Exposes a number of signals so the view can be monitored7- Tracks focus within the view8- Exposes a settings store for flows that automatically expires if the flow is9  removed from the store.10"""11 12import collections13import logging14import re15from collections.abc import Iterator16from collections.abc import MutableMapping17from collections.abc import Sequence18from typing import Any19from typing import Optional20 21import sortedcontainers22 23import mitmproxy.flow24from mitmproxy import command25from mitmproxy import connection26from mitmproxy import ctx27from mitmproxy import dns28from mitmproxy import exceptions29from mitmproxy import flowfilter30from mitmproxy import hooks31from mitmproxy import http32from mitmproxy import io33from mitmproxy import tcp34from mitmproxy import udp35from mitmproxy.log import ALERT36from mitmproxy.utils import human37from mitmproxy.utils import signals38 39# The underlying sorted list implementation expects the sort key to be stable40# for the lifetime of the object. However, if we sort by size, for instance,41# the sort order changes as the flow progresses through its lifecycle. We42# address this through two means:43#44# - Let order keys cache the sort value by flow ID.45#46# - Add a facility to refresh items in the list by removing and re-adding them47# when they are updated.48 49 50class _OrderKey:51    def __init__(self, view):52        self.view = view53 54    def generate(self, f: mitmproxy.flow.Flow) -> Any:  # pragma: no cover55        pass56 57    def refresh(self, f):58        k = self._key()59        old = self.view.settings[f][k]60        new = self.generate(f)61        if old != new:62            self.view._view.remove(f)63            self.view.settings[f][k] = new64            self.view._view.add(f)65            self.view.sig_view_refresh.send()66 67    def _key(self):68        return "_order_%s" % id(self)69 70    def __call__(self, f):71        if f.id in self.view._store:72            k = self._key()73            s = self.view.settings[f]74            if k in s:75                return s[k]76            val = self.generate(f)77            s[k] = val78            return val79        else:80            return self.generate(f)81 82 83class OrderRequestStart(_OrderKey):84    def generate(self, f: mitmproxy.flow.Flow) -> float:85        return f.timestamp_created86 87 88class OrderRequestMethod(_OrderKey):89    def generate(self, f: mitmproxy.flow.Flow) -> str:90        if isinstance(f, http.HTTPFlow):91            return f.request.method92        elif isinstance(f, (tcp.TCPFlow, udp.UDPFlow)):93            return f.type.upper()94        elif isinstance(f, dns.DNSFlow):95            return dns.op_codes.to_str(f.request.op_code)96        else:97            raise NotImplementedError()98 99 100class OrderRequestURL(_OrderKey):101    def generate(self, f: mitmproxy.flow.Flow) -> str:102        if isinstance(f, http.HTTPFlow):103            return f.request.url104        elif isinstance(f, (tcp.TCPFlow, udp.UDPFlow)):105            return human.format_address(f.server_conn.address)106        elif isinstance(f, dns.DNSFlow):107            return f.request.questions[0].name if f.request.questions else ""108        else:109            raise NotImplementedError()110 111 112class OrderKeySize(_OrderKey):113    def generate(self, f: mitmproxy.flow.Flow) -> int:114        if isinstance(f, http.HTTPFlow):115            size = 0116            if f.request.raw_content:117                size += len(f.request.raw_content)118            if f.response and f.response.raw_content:119                size += len(f.response.raw_content)120            return size121        elif isinstance(f, (tcp.TCPFlow, udp.UDPFlow)):122            size = 0123            for message in f.messages:124                size += len(message.content)125            return size126        elif isinstance(f, dns.DNSFlow):127            return f.response.size if f.response else 0128        else:129            raise NotImplementedError()130 131 132orders = [133    ("t", "time"),134    ("m", "method"),135    ("u", "url"),136    ("z", "size"),137]138 139 140def _signal_with_flow(flow: mitmproxy.flow.Flow) -> None: ...141 142 143def _sig_view_remove(flow: mitmproxy.flow.Flow, index: int) -> None: ...144 145 146class View(collections.abc.Sequence):147    def __init__(self) -> None:148        super().__init__()149        self._store: collections.OrderedDict[str, mitmproxy.flow.Flow] = (150            collections.OrderedDict()151        )152        self.filter = flowfilter.match_all153        # Should we show only marked flows?154        self.show_marked = False155 156        self.default_order = OrderRequestStart(self)157        self.orders = dict(158            time=OrderRequestStart(self),159            method=OrderRequestMethod(self),160            url=OrderRequestURL(self),161            size=OrderKeySize(self),162        )163        self.order_key: _OrderKey = self.default_order164        self.order_reversed = False165        self.focus_follow = False166 167        self._view = sortedcontainers.SortedListWithKey(key=self.order_key)168 169        # The sig_view* signals broadcast events that affect the view. That is,170        # an update to a flow in the store but not in the view does not trigger171        # a signal. All signals are called after the view has been updated.172        self.sig_view_update = signals.SyncSignal(_signal_with_flow)173        self.sig_view_add = signals.SyncSignal(_signal_with_flow)174        self.sig_view_remove = signals.SyncSignal(_sig_view_remove)175        # Signals that the view should be refreshed completely176        self.sig_view_refresh = signals.SyncSignal(lambda: None)177 178        # The sig_store* signals broadcast events that affect the underlying179        # store. If a flow is removed from just the view, sig_view_remove is180        # triggered. If it is removed from the store while it is also in the181        # view, both sig_store_remove and sig_view_remove are triggered.182        self.sig_store_remove = signals.SyncSignal(_signal_with_flow)183        # Signals that the store should be refreshed completely184        self.sig_store_refresh = signals.SyncSignal(lambda: None)185 186        self.focus = Focus(self)187        self.settings = Settings(self)188 189    def load(self, loader):190        loader.add_option(191            "view_filter", Optional[str], None, "Limit the view to matching flows."192        )193        loader.add_option(194            "view_order",195            str,196            "time",197            "Flow sort order.",198            choices=list(map(lambda c: c[1], orders)),199        )200        loader.add_option(201            "view_order_reversed", bool, False, "Reverse the sorting order."202        )203        loader.add_option(204            "console_focus_follow", bool, False, "Focus follows new flows."205        )206 207    def store_count(self):208        return len(self._store)209 210    def _rev(self, idx: int) -> int:211        """212        Reverses an index, if needed213        """214        if self.order_reversed:215            if idx < 0:216                idx = -idx - 1217            else:218                idx = len(self._view) - idx - 1219                if idx < 0:220                    raise IndexError221        return idx222 223    def __len__(self):224        return len(self._view)225 226    def __getitem__(self, offset) -> Any:227        return self._view[self._rev(offset)]228 229    # Reflect some methods to the efficient underlying implementation230 231    def _bisect(self, f: mitmproxy.flow.Flow) -> int:232        v = self._view.bisect_right(f)233        return self._rev(v - 1) + 1234 235    def index(236        self, f: mitmproxy.flow.Flow, start: int = 0, stop: int | None = None237    ) -> int:238        return self._rev(self._view.index(f, start, stop))239 240    def __contains__(self, f: Any) -> bool:241        return self._view.__contains__(f)242 243    def _order_key_name(self):244        return "_order_%s" % id(self.order_key)245 246    def _base_add(self, f):247        self.settings[f][self._order_key_name()] = self.order_key(f)248        self._view.add(f)249 250    def _refilter(self):251        self._view.clear()252        for i in self._store.values():253            if self.show_marked and not i.marked:254                continue255            if self.filter(i):256                self._base_add(i)257        self.sig_view_refresh.send()258 259    """ View API """260 261    # Focus262    @command.command("view.focus.go")263    def go(self, offset: int) -> None:264        """265        Go to a specified offset. Positive offests are from the beginning of266        the view, negative from the end of the view, so that 0 is the first267        flow, -1 is the last flow.268        """269        if len(self) == 0:270            return271        if offset < 0:272            offset = len(self) + offset273        if offset < 0:274            offset = 0275        if offset > len(self) - 1:276            offset = len(self) - 1277        self.focus.flow = self[offset]278 279    @command.command("view.focus.next")280    def focus_next(self) -> None:281        """282        Set focus to the next flow.283        """284        if self.focus.index is not None:285            idx = self.focus.index + 1286            if self.inbounds(idx):287                self.focus.flow = self[idx]288        else:289            pass290 291    @command.command("view.focus.prev")292    def focus_prev(self) -> None:293        """294        Set focus to the previous flow.295        """296        if self.focus.index is not None:297            idx = self.focus.index - 1298            if self.inbounds(idx):299                self.focus.flow = self[idx]300        else:301            pass302 303    # Order304    @command.command("view.order.options")305    def order_options(self) -> Sequence[str]:306        """307        Choices supported by the view_order option.308        """309        return list(sorted(self.orders.keys()))310 311    @command.command("view.order.reverse")312    def set_reversed(self, boolean: bool) -> None:313        self.order_reversed = boolean314        self.sig_view_refresh.send()315 316    @command.command("view.order.set")317    def set_order(self, order_key: str) -> None:318        """319        Sets the current view order.320        """321        if order_key not in self.orders:322            raise exceptions.CommandError("Unknown flow order: %s" % order_key)323        key = self.orders[order_key]324        self.order_key = key325        newview = sortedcontainers.SortedListWithKey(key=key)326        newview.update(self._view)327        self._view = newview328 329    @command.command("view.order")330    def get_order(self) -> str:331        """332        Returns the current view order.333        """334        order = ""335        for k in self.orders.keys():336            if self.order_key == self.orders[k]:337                order = k338        return order339 340    # Filter341    @command.command("view.filter.set")342    def set_filter_cmd(self, filter_expr: str) -> None:343        """344        Sets the current view filter.345        """346        filt = None347        if filter_expr:348            try:349                filt = flowfilter.parse(filter_expr)350            except ValueError as e:351                raise exceptions.CommandError(str(e)) from e352        self.set_filter(filt)353 354    def set_filter(self, flt: flowfilter.TFilter | None):355        self.filter = flt or flowfilter.match_all356        self._refilter()357 358    # View Updates359    @command.command("view.clear")360    def clear(self) -> None:361        """362        Clears both the store and view.363        """364        self._store.clear()365        self._view.clear()366        self.sig_view_refresh.send()367        self.sig_store_refresh.send()368 369    @command.command("view.clear_unmarked")370    def clear_not_marked(self) -> None:371        """372        Clears only the unmarked flows.373        """374        for flow in self._store.copy().values():375            if not flow.marked:376                self._store.pop(flow.id)377 378        self._refilter()379        self.sig_store_refresh.send()380 381    # View Settings382    @command.command("view.settings.getval")383    def getvalue(self, flow: mitmproxy.flow.Flow, key: str, default: str) -> str:384        """385        Get a value from the settings store for the specified flow.386        """387        return self.settings[flow].get(key, default)388 389    @command.command("view.settings.setval.toggle")390    def setvalue_toggle(self, flows: Sequence[mitmproxy.flow.Flow], key: str) -> None:391        """392        Toggle a boolean value in the settings store, setting the value to393        the string "true" or "false".394        """395        updated = []396        for f in flows:397            current = self.settings[f].get(key, "false")398            self.settings[f][key] = "false" if current == "true" else "true"399            updated.append(f)400        ctx.master.addons.trigger(hooks.UpdateHook(updated))401 402    @command.command("view.settings.setval")403    def setvalue(404        self, flows: Sequence[mitmproxy.flow.Flow], key: str, value: str405    ) -> None:406        """407        Set a value in the settings store for the specified flows.408        """409        updated = []410        for f in flows:411            self.settings[f][key] = value412            updated.append(f)413        ctx.master.addons.trigger(hooks.UpdateHook(updated))414 415    # Flows416    @command.command("view.flows.duplicate")417    def duplicate(self, flows: Sequence[mitmproxy.flow.Flow]) -> None:418        """419        Duplicates the specified flows, and sets the focus to the first420        duplicate.421        """422        dups = [f.copy() for f in flows]423        if dups:424            self.add(dups)425            self.focus.flow = dups[0]426            logging.log(ALERT, "Duplicated %s flows" % len(dups))427 428    @command.command("view.flows.remove")429    def remove(self, flows: Sequence[mitmproxy.flow.Flow]) -> None:430        """431        Removes the flow from the underlying store and the view.432        """433        for f in flows:434            if f.id in self._store:435                if f.killable:436                    f.kill()437                if f in self._view:438                    # We manually pass the index here because multiple flows may have the same439                    # sorting key, and we cannot reconstruct the index from that.440                    idx = self._view.index(f)441                    self._view.remove(f)442                    self.sig_view_remove.send(flow=f, index=idx)443                del self._store[f.id]444                self.sig_store_remove.send(flow=f)445        if len(flows) > 1:446            logging.log(ALERT, "Removed %s flows" % len(flows))447 448    @command.command("view.flows.resolve")449    def resolve(self, flow_spec: str) -> Sequence[mitmproxy.flow.Flow]:450        """451        Resolve a flow list specification to an actual list of flows.452        """453        if flow_spec == "@all":454            return [i for i in self._store.values()]455        if flow_spec == "@focus":456            return [self.focus.flow] if self.focus.flow else []457        elif flow_spec == "@shown":458            return [i for i in self]459        elif flow_spec == "@hidden":460            return [i for i in self._store.values() if i not in self._view]461        elif flow_spec == "@marked":462            return [i for i in self._store.values() if i.marked]463        elif flow_spec == "@unmarked":464            return [i for i in self._store.values() if not i.marked]465        elif re.match(r"@[0-9a-f\-,]{36,}", flow_spec):466            ids = flow_spec[1:].split(",")467            return [i for i in self._store.values() if i.id in ids]468        else:469            try:470                filt = flowfilter.parse(flow_spec)471            except ValueError as e:472                raise exceptions.CommandError(str(e)) from e473            return [i for i in self._store.values() if filt(i)]474 475    @command.command("view.flows.create")476    def create(self, method: str, url: str) -> None:477        try:478            req = http.Request.make(method.upper(), url)479        except ValueError as e:480            raise exceptions.CommandError("Invalid URL: %s" % e)481 482        c = connection.Client(483            peername=("", 0),484            sockname=("", 0),485            timestamp_start=req.timestamp_start - 0.0001,486        )487        s = connection.Server(address=(req.host, req.port))488 489        f = http.HTTPFlow(c, s)490        f.request = req491        f.request.headers["Host"] = req.host492        self.add([f])493 494    @command.command("view.flows.load")495    def load_file(self, path: mitmproxy.types.Path) -> None:496        """497        Load flows into the view, without processing them with addons.498        """499        try:500            with open(path, "rb") as f:501                for i in io.FlowReader(f).stream():502                    # Do this to get a new ID, so we can load the same file N times and503                    # get new flows each time. It would be more efficient to just have a504                    # .newid() method or something.505                    self.add([i.copy()])506        except OSError as e:507            logging.error(e.strerror)508        except exceptions.FlowReadException as e:509            logging.error(str(e))510 511    def add(self, flows: Sequence[mitmproxy.flow.Flow]) -> None:512        """513        Adds a flow to the state. If the flow already exists, it is514        ignored.515        """516        for f in flows:517            if f.id not in self._store:518                self._store[f.id] = f519                if self.filter(f):520                    self._base_add(f)521                    if self.focus_follow:522                        self.focus.flow = f523                    self.sig_view_add.send(flow=f)524 525    def get_by_id(self, flow_id: str) -> mitmproxy.flow.Flow | None:526        """527        Get flow with the given id from the store.528        Returns None if the flow is not found.529        """530        return self._store.get(flow_id)531 532    # View Properties533    @command.command("view.properties.length")534    def get_length(self) -> int:535        """536        Returns view length.537        """538        return len(self)539 540    @command.command("view.properties.marked")541    def get_marked(self) -> bool:542        """543        Returns true if view is in marked mode.544        """545        return self.show_marked546 547    @command.command("view.properties.marked.toggle")548    def toggle_marked(self) -> None:549        """550        Toggle whether to show marked views only.551        """552        self.show_marked = not self.show_marked553        self._refilter()554 555    @command.command("view.properties.inbounds")556    def inbounds(self, index: int) -> bool:557        """558        Is this 0 <= index < len(self)?559        """560        return 0 <= index < len(self)561 562    # Event handlers563    def configure(self, updated):564        if "view_filter" in updated:565            filt = None566            if ctx.options.view_filter:567                try:568                    filt = flowfilter.parse(ctx.options.view_filter)569                except ValueError as e:570                    raise exceptions.OptionsError(str(e)) from e571            self.set_filter(filt)572        if "view_order" in updated:573            if ctx.options.view_order not in self.orders:574                raise exceptions.OptionsError(575                    "Unknown flow order: %s" % ctx.options.view_order576                )577            self.set_order(ctx.options.view_order)578        if "view_order_reversed" in updated:579            self.set_reversed(ctx.options.view_order_reversed)580        if "console_focus_follow" in updated:581            self.focus_follow = ctx.options.console_focus_follow582 583    def requestheaders(self, f):584        self.add([f])585 586    def error(self, f):587        self.update([f])588 589    def response(self, f):590        self.update([f])591 592    def intercept(self, f):593        self.update([f])594 595    def resume(self, f):596        self.update([f])597 598    def kill(self, f):599        self.update([f])600 601    def tcp_start(self, f):602        self.add([f])603 604    def tcp_message(self, f):605        self.update([f])606 607    def tcp_error(self, f):608        self.update([f])609 610    def tcp_end(self, f):611        self.update([f])612 613    def udp_start(self, f):614        self.add([f])615 616    def udp_message(self, f):617        self.update([f])618 619    def udp_error(self, f):620        self.update([f])621 622    def udp_end(self, f):623        self.update([f])624 625    def dns_request(self, f):626        self.add([f])627 628    def dns_response(self, f):629        self.update([f])630 631    def dns_error(self, f):632        self.update([f])633 634    def update(self, flows: Sequence[mitmproxy.flow.Flow]) -> None:635        """636        Updates a list of flows. If flow is not in the state, it's ignored.637        """638        for f in flows:639            if f.id in self._store:640                if self.filter(f):641                    if f not in self._view:642                        self._base_add(f)643                        if self.focus_follow:644                            self.focus.flow = f645                        self.sig_view_add.send(flow=f)646                    else:647                        # This is a tad complicated. The sortedcontainers648                        # implementation assumes that the order key is stable. If649                        # it changes mid-way Very Bad Things happen. We detect when650                        # this happens, and re-fresh the item.651                        self.order_key.refresh(f)652                        self.sig_view_update.send(flow=f)653                else:654                    try:655                        idx = self._view.index(f)656                    except ValueError:657                        pass  # The value was not in the view658                    else:659                        self._view.remove(f)660                        self.sig_view_remove.send(flow=f, index=idx)661 662 663class Focus:664    """665    Tracks a focus element within a View.666    """667 668    def __init__(self, v: View) -> None:669        self.view = v670        self._flow: mitmproxy.flow.Flow | None = None671        self.sig_change = signals.SyncSignal(lambda: None)672        if len(self.view):673            self.flow = self.view[0]674        v.sig_view_add.connect(self._sig_view_add)675        v.sig_view_remove.connect(self._sig_view_remove)676        v.sig_view_refresh.connect(self._sig_view_refresh)677 678    @property679    def flow(self) -> mitmproxy.flow.Flow | None:680        return self._flow681 682    @flow.setter683    def flow(self, f: mitmproxy.flow.Flow | None):684        if f is not None and f not in self.view:685            raise ValueError("Attempt to set focus to flow not in view")686        self._flow = f687        self.sig_change.send()688 689    @property690    def index(self) -> int | None:691        if self.flow:692            return self.view.index(self.flow)693        return None694 695    @index.setter696    def index(self, idx):697        if idx < 0 or idx > len(self.view) - 1:698            raise ValueError("Index out of view bounds")699        self.flow = self.view[idx]700 701    def _nearest(self, f, v):702        return min(v._bisect(f), len(v) - 1)703 704    def _sig_view_remove(self, flow, index):705        if len(self.view) == 0:706            self.flow = None707        elif flow is self.flow:708            self.index = min(index, len(self.view) - 1)709 710    def _sig_view_refresh(self):711        if len(self.view) == 0:712            self.flow = None713        elif self.flow is None:714            self.flow = self.view[0]715        elif self.flow not in self.view:716            self.flow = self.view[self._nearest(self.flow, self.view)]717 718    def _sig_view_add(self, flow):719        # We only have to act if we don't have a focus element720        if not self.flow:721            self.flow = flow722 723 724class Settings(collections.abc.Mapping):725    def __init__(self, view: View) -> None:726        self.view = view727        self._values: MutableMapping[str, dict] = {}728        view.sig_store_remove.connect(self._sig_store_remove)729        view.sig_store_refresh.connect(self._sig_store_refresh)730 731    def __iter__(self) -> Iterator:732        return iter(self._values)733 734    def __len__(self) -> int:735        return len(self._values)736 737    def __getitem__(self, f: mitmproxy.flow.Flow) -> dict:738        if f.id not in self.view._store:739            raise KeyError740        return self._values.setdefault(f.id, {})741 742    def _sig_store_remove(self, flow):743        if flow.id in self._values:744            del self._values[flow.id]745 746    def _sig_store_refresh(self):747        for fid in list(self._values.keys()):748            if fid not in self.view._store:749                del self._values[fid]750