codekingpro/portable-devtools

1import logging2import os.path3import sys4from collections.abc import Sequence5from datetime import datetime6from functools import lru_cache7from pathlib import Path8from typing import Literal9from typing import Optional10 11import mitmproxy.types12from mitmproxy import command13from mitmproxy import ctx14from mitmproxy import dns15from mitmproxy import exceptions16from mitmproxy import flow17from mitmproxy import flowfilter18from mitmproxy import http19from mitmproxy import io20from mitmproxy import tcp21from mitmproxy import udp22from mitmproxy.log import ALERT23 24 25@lru_cache26def _path(path: str) -> str:27    """Extract the path from a path spec (which may have an extra "+" at the front)"""28    if path.startswith("+"):29        path = path[1:]30    return os.path.expanduser(path)31 32 33@lru_cache34def _mode(path: str) -> Literal["ab", "wb"]:35    """Extract the writing mode (overwrite or append) from a path spec"""36    if path.startswith("+"):37        return "ab"38    else:39        return "wb"40 41 42class Save:43    def __init__(self) -> None:44        self.stream: io.FilteredFlowWriter | None = None45        self.filt: flowfilter.TFilter | None = None46        self.active_flows: set[flow.Flow] = set()47        self.current_path: str | None = None48 49    def load(self, loader):50        loader.add_option(51            "save_stream_file",52            Optional[str],53            None,54            """55            Stream flows to file as they arrive. Prefix path with + to append.56            The full path can use python strftime() formating, missing57            directories are created as needed. A new file is opened every time58            the formatted string changes.59            """,60        )61        loader.add_option(62            "save_stream_filter",63            Optional[str],64            None,65            "Filter which flows are written to file.",66        )67 68    def configure(self, updated):69        if "save_stream_filter" in updated:70            if ctx.options.save_stream_filter:71                try:72                    self.filt = flowfilter.parse(ctx.options.save_stream_filter)73                except ValueError as e:74                    raise exceptions.OptionsError(str(e)) from e75            else:76                self.filt = None77        if "save_stream_file" in updated or "save_stream_filter" in updated:78            if ctx.options.save_stream_file:79                try:80                    self.maybe_rotate_to_new_file()81                except OSError as e:82                    raise exceptions.OptionsError(str(e)) from e83                assert self.stream84                self.stream.flt = self.filt85            else:86                self.done()87 88    def maybe_rotate_to_new_file(self) -> None:89        path = datetime.today().strftime(_path(ctx.options.save_stream_file))90        if self.current_path == path:91            return92 93        if self.stream:94            self.stream.fo.close()95            self.stream = None96 97        new_log_file = Path(path)98        new_log_file.parent.mkdir(parents=True, exist_ok=True)99 100        f = new_log_file.open(_mode(ctx.options.save_stream_file))101        self.stream = io.FilteredFlowWriter(f, self.filt)102        self.current_path = path103 104    def save_flow(self, flow: flow.Flow) -> None:105        """106        Write the flow to the stream, but first check if we need to rotate to a new file.107        """108        if not self.stream:109            return110        try:111            self.maybe_rotate_to_new_file()112            self.stream.add(flow)113        except OSError as e:114            # If we somehow fail to write flows to a logfile, we really want to crash visibly115            # instead of letting traffic through unrecorded.116            # No normal logging here, that would not be triggered anymore.117            sys.stderr.write(f"Error while writing to {self.current_path}: {e}")118            sys.exit(1)119        else:120            self.active_flows.discard(flow)121 122    def done(self) -> None:123        if self.stream:124            for f in self.active_flows:125                self.stream.add(f)126            self.active_flows.clear()127 128            self.current_path = None129            self.stream.fo.close()130            self.stream = None131 132    @command.command("save.file")133    def save(self, flows: Sequence[flow.Flow], path: mitmproxy.types.Path) -> None:134        """135        Save flows to a file. If the path starts with a +, flows are136        appended to the file, otherwise it is over-written.137        """138        try:139            with open(_path(path), _mode(path)) as f:140                stream = io.FlowWriter(f)141                for i in flows:142                    stream.add(i)143        except OSError as e:144            raise exceptions.CommandError(e) from e145        if path.endswith(".har") or path.endswith(".zhar"):  # pragma: no cover146            logging.log(147                ALERT,148                f"Saved as mitmproxy dump file. To save HAR files, use the `save.har` command.",149            )150        else:151            logging.log(ALERT, f"Saved {len(flows)} flows.")152 153    def tcp_start(self, flow: tcp.TCPFlow):154        if self.stream:155            self.active_flows.add(flow)156 157    def tcp_end(self, flow: tcp.TCPFlow):158        self.save_flow(flow)159 160    def tcp_error(self, flow: tcp.TCPFlow):161        self.tcp_end(flow)162 163    def udp_start(self, flow: udp.UDPFlow):164        if self.stream:165            self.active_flows.add(flow)166 167    def udp_end(self, flow: udp.UDPFlow):168        self.save_flow(flow)169 170    def udp_error(self, flow: udp.UDPFlow):171        self.udp_end(flow)172 173    def websocket_end(self, flow: http.HTTPFlow):174        self.save_flow(flow)175 176    def request(self, flow: http.HTTPFlow):177        if self.stream:178            self.active_flows.add(flow)179 180    def response(self, flow: http.HTTPFlow):181        # websocket flows will receive a websocket_end,182        # we don't want to persist them here already183        if flow.websocket is None:184            self.save_flow(flow)185 186    def error(self, flow: http.HTTPFlow):187        self.response(flow)188 189    def dns_request(self, flow: dns.DNSFlow):190        if self.stream:191            self.active_flows.add(flow)192 193    def dns_response(self, flow: dns.DNSFlow):194        self.save_flow(flow)195 196    def dns_error(self, flow: dns.DNSFlow):197        self.save_flow(flow)198