mitmproxy tcp-script

tcp-save.py

import os.path
import typing

from mitmproxy import command
from mitmproxy import flow
from mitmproxy import tcp
from mitmproxy import exceptions
from mitmproxy import ctx
import json
import base64
import mitmproxy.types


class TcpSave:
    def open_file(self, path):
        path = os.path.expanduser(path)
        return open(path, "wb")

    @command.command("tcp.save")
    def save(self,
             flows: typing.Sequence[flow.Flow],
             path: mitmproxy.types.Path) -> None:
        """
            Save tcp-flows to a json file.
        """

        for flow in flows:
            if not isinstance(flow, tcp.TCPFlow):
                raise exceptions.CommandError("Only TCPFlows can be saved.")

        try:
            f = self.open_file(path)
        except OSError as v:
            raise exceptions.CommandError(v) from v

        b64_messages = list(map(lambda flow: 
            list(map(lambda chunk: 
                base64.b64encode(chunk.content).decode('utf-8'),
                flow.messages)
            ), flows)
        )
        f.write(json.dumps(b64_messages).encode('utf-8'))
        f.close()
        ctx.log.alert(f"Saved {len(flows)} flows.")


addons = [
    TcpSave()
]