codekingpro/portable-devtools

1import json2import os3from collections.abc import Iterable4from io import BufferedReader5from typing import Any6from typing import BinaryIO7from typing import cast8from typing import Union9 10from mitmproxy import exceptions11from mitmproxy import flow12from mitmproxy import flowfilter13from mitmproxy.io import compat14from mitmproxy.io import tnetstring15from mitmproxy.io.har import request_to_flow16 17 18class FlowWriter:19    def __init__(self, fo):20        self.fo = fo21 22    def add(self, f: flow.Flow) -> None:23        d = f.get_state()24        tnetstring.dump(d, self.fo)25 26 27class FlowReader:28    fo: BinaryIO29 30    def __init__(self, fo: BinaryIO):31        self.fo = fo32 33    def peek(self, n: int) -> bytes:34        try:35            return cast(BufferedReader, self.fo).peek(n)36        except AttributeError:37            # https://github.com/python/cpython/issues/90533: io.BytesIO does not have peek()38            pos = self.fo.tell()39            ret = self.fo.read(n)40            self.fo.seek(pos)41            return ret42 43    def stream(self) -> Iterable[flow.Flow]:44        """45        Yields Flow objects from the dump.46        """47 48        if self.peek(4).startswith(49            b"\xef\xbb\xbf{"50        ):  # skip BOM, usually added by Fiddler51            self.fo.read(3)52        if self.peek(1).startswith(b"{"):53            try:54                har_file = json.loads(self.fo.read().decode("utf-8"))55 56                for request_json in har_file["log"]["entries"]:57                    yield request_to_flow(request_json)58 59            except Exception:60                raise exceptions.FlowReadException(61                    "Unable to read HAR file. Please provide a valid HAR file"62                )63 64        else:65            try:66                while True:67                    # FIXME: This cast hides a lack of dynamic type checking68                    loaded = cast(69                        dict[Union[bytes, str], Any],70                        tnetstring.load(self.fo),71                    )72                    try:73                        if not isinstance(loaded, dict):74                            raise ValueError(f"Invalid flow: {loaded=}")75                        yield flow.Flow.from_state(compat.migrate_flow(loaded))76                    except ValueError as e:77                        raise exceptions.FlowReadException(e) from e78            except (ValueError, TypeError, IndexError) as e:79                if str(e) == "not a tnetstring: empty file":80                    return  # Error is due to EOF81                raise exceptions.FlowReadException("Invalid data format.") from e82 83 84class FilteredFlowWriter:85    def __init__(self, fo: BinaryIO, flt: flowfilter.TFilter | None):86        self.fo = fo87        self.flt = flt88 89    def add(self, f: flow.Flow) -> None:90        if self.flt and not flowfilter.match(self.flt, f):91            return92        d = f.get_state()93        tnetstring.dump(d, self.fo)94        self.fo.flush()95 96 97def read_flows_from_paths(paths) -> list[flow.Flow]:98    """99    Given a list of filepaths, read all flows and return a list of them.100    From a performance perspective, streaming would be advisable -101    however, if there's an error with one of the files, we want it to be raised immediately.102 103    Raises:104        FlowReadException, if any error occurs.105    """106    try:107        flows: list[flow.Flow] = []108        for path in paths:109            path = os.path.expanduser(path)110            with open(path, "rb") as f:111                flows.extend(FlowReader(f).stream())112    except OSError as e:113        raise exceptions.FlowReadException(e.strerror)114    return flows115