diff options
Diffstat (limited to 'mitmproxy/addons/serverplayback.py')
| -rw-r--r-- | mitmproxy/addons/serverplayback.py | 67 |
1 files changed, 46 insertions, 21 deletions
diff --git a/mitmproxy/addons/serverplayback.py b/mitmproxy/addons/serverplayback.py index 51ba60b4..7f642585 100644 --- a/mitmproxy/addons/serverplayback.py +++ b/mitmproxy/addons/serverplayback.py @@ -1,16 +1,19 @@ import hashlib -import urllib import typing +import urllib -from mitmproxy import ctx -from mitmproxy import flow +import mitmproxy.types +from mitmproxy import command +from mitmproxy import ctx, http from mitmproxy import exceptions +from mitmproxy import flow from mitmproxy import io -from mitmproxy import command -import mitmproxy.types class ServerPlayback: + flowmap: typing.Dict[typing.Hashable, typing.List[http.HTTPFlow]] + configured: bool + def __init__(self): self.flowmap = {} self.configured = False @@ -68,6 +71,13 @@ class ServerPlayback: to replay. """ ) + loader.add_option( + "server_replay_ignore_port", bool, False, + """ + Ignore request's destination port while searching for a saved flow + to replay. + """ + ) @command.command("replay.server") def load_flows(self, flows: typing.Sequence[flow.Flow]) -> None: @@ -75,10 +85,10 @@ class ServerPlayback: Replay server responses from flows. """ self.flowmap = {} - for i in flows: - if i.response: # type: ignore - l = self.flowmap.setdefault(self._hash(i), []) - l.append(i) + for f in flows: + if isinstance(f, http.HTTPFlow): + lst = self.flowmap.setdefault(self._hash(f), []) + lst.append(f) ctx.master.addons.trigger("update", []) @command.command("replay.server.file") @@ -101,16 +111,15 @@ class ServerPlayback: def count(self) -> int: return sum([len(i) for i in self.flowmap.values()]) - def _hash(self, flow): + def _hash(self, flow: http.HTTPFlow) -> typing.Hashable: """ 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: typing.List[typing.Any] = [str(r.port), str(r.scheme), str(r.method), str(path)] + key: typing.List[typing.Any] = [str(r.scheme), str(r.method), str(path)] if not ctx.options.server_replay_ignore_content: if ctx.options.server_replay_ignore_payload_params and r.multipart_form: key.extend( @@ -128,7 +137,9 @@ class ServerPlayback: key.append(str(r.raw_content)) if not ctx.options.server_replay_ignore_host: - key.append(r.host) + key.append(r.pretty_host) + if not ctx.options.server_replay_ignore_port: + key.append(r.port) filtered = [] ignore_params = ctx.options.server_replay_ignore_params or [] @@ -149,20 +160,32 @@ class ServerPlayback: repr(key).encode("utf8", "surrogateescape") ).digest() - def next_flow(self, request): + def next_flow(self, flow: http.HTTPFlow) -> typing.Optional[http.HTTPFlow]: """ Returns the next flow object, or None if no matching flow was found. """ - hsh = self._hash(request) - if hsh in self.flowmap: + hash = self._hash(flow) + if hash in self.flowmap: if ctx.options.server_replay_nopop: - return self.flowmap[hsh][0] + return next(( + flow + for flow in self.flowmap[hash] + if flow.response + ), None) else: - ret = self.flowmap[hsh].pop(0) - if not self.flowmap[hsh]: - del self.flowmap[hsh] + ret = self.flowmap[hash].pop(0) + while not ret.response: + if self.flowmap[hash]: + ret = self.flowmap[hash].pop(0) + else: + del self.flowmap[hash] + return None + if not self.flowmap[hash]: + del self.flowmap[hash] return ret + else: + return None def configure(self, updated): if not self.configured and ctx.options.server_replay: @@ -173,10 +196,11 @@ class ServerPlayback: raise exceptions.OptionsError(str(e)) self.load_flows(flows) - def request(self, f): + def request(self, f: http.HTTPFlow) -> None: if self.flowmap: rflow = self.next_flow(f) if rflow: + assert rflow.response response = rflow.response.copy() response.is_replay = True if ctx.options.server_replay_refresh: @@ -188,4 +212,5 @@ class ServerPlayback: f.request.url ) ) + assert f.reply f.reply.kill() |