aboutsummaryrefslogtreecommitdiffstats
path: root/mitmproxy/addons/serverplayback.py
blob: 1161ce23414a3c71d29df6706c679fbc4eb2b6af (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
import urllib
import hashlib

from netlib import strutils
from mitmproxy import exceptions
from mitmproxy import ctx
from mitmproxy import io


class ServerPlayback:
    def __init__(self):
        self.options = None

        self.flowmap = {}
        self.stop = False
        self.final_flow = None

    def load(self, flows):
        for i in flows:
            if i.response:
                l = self.flowmap.setdefault(self._hash(i), [])
                l.append(i)

    def clear(self):
        self.flowmap = {}

    def count(self):
        return sum([len(i) for i in self.flowmap.values()])

    def _hash(self, flow):
        """
            Calculates a loose hash of the flow request.
        """
        r = flow.request

        _, _, path, _, query, _ = urllib.parse.urlparse(r.url)
        queriesArray = urllib.parse.parse_qsl(query, keep_blank_values=True)

        key = [str(r.port), str(r.scheme), str(r.method), str(path)]
        if not self.options.server_replay_ignore_content:
            form_contents = r.urlencoded_form or r.multipart_form
            if self.options.server_replay_ignore_payload_params and form_contents:
                params = [
                    strutils.always_bytes(i)
                    for i in self.options.server_replay_ignore_payload_params
                ]
                for p in form_contents.items(multi=True):
                    if p[0] not in params:
                        key.append(p)
            else:
                key.append(str(r.raw_content))

        if not self.options.server_replay_ignore_host:
            key.append(r.host)

        filtered = []
        ignore_params = self.options.server_replay_ignore_params or []
        for p in queriesArray:
            if p[0] not in ignore_params:
                filtered.append(p)
        for p in filtered:
            key.append(p[0])
            key.append(p[1])

        if self.options.server_replay_use_headers:
            headers = []
            for i in self.options.server_replay_use_headers:
                v = r.headers.get(i)
                headers.append((i, v))
            key.append(headers)
        return hashlib.sha256(
            repr(key).encode("utf8", "surrogateescape")
        ).digest()

    def next_flow(self, request):
        """
            Returns the next flow object, or None if no matching flow was
            found.
        """
        hsh = self._hash(request)
        if hsh in self.flowmap:
            if self.options.server_replay_nopop:
                return self.flowmap[hsh][0]
            else:
                ret = self.flowmap[hsh].pop(0)
                if not self.flowmap[hsh]:
                    del self.flowmap[hsh]
                return ret

    def configure(self, options, updated):
        self.options = options
        if "server_replay" in updated:
            self.clear()
            if options.server_replay:
                try:
                    flows = io.read_flows_from_paths(options.server_replay)
                except exceptions.FlowReadException as e:
                    raise exceptions.OptionsError(str(e))
                self.load(flows)

    def tick(self):
        if self.stop and not self.final_flow.live:
            ctx.master.shutdown()

    def request(self, f):
        if self.flowmap:
            rflow = self.next_flow(f)
            if rflow:
                response = rflow.response.copy()
                response.is_replay = True
                if self.options.refresh_server_playback:
                    response.refresh()
                f.response = response
                if not self.flowmap and not self.options.keepserving:
                    self.final_flow = f
                    self.stop = True
            elif self.options.replay_kill_extra:
                ctx.log.warn(
                    "server_playback: killed non-replay request {}".format(
                        f.request.url
                    )
                )
                f.reply.kill()