diff options
Diffstat (limited to 'mitmproxy/addons/serverplayback.py')
| -rw-r--r-- | mitmproxy/addons/serverplayback.py | 75 |
1 files changed, 47 insertions, 28 deletions
diff --git a/mitmproxy/addons/serverplayback.py b/mitmproxy/addons/serverplayback.py index be2d6f2b..927f6e15 100644 --- a/mitmproxy/addons/serverplayback.py +++ b/mitmproxy/addons/serverplayback.py @@ -1,29 +1,50 @@ import hashlib import urllib +import typing from typing import Any # noqa from typing import List # noqa from mitmproxy import ctx +from mitmproxy import flow from mitmproxy import exceptions from mitmproxy import io +from mitmproxy import command class ServerPlayback: def __init__(self): - self.options = None - self.flowmap = {} self.stop = False self.final_flow = None + self.configured = False - def load(self, flows): + @command.command("replay.server") + def load_flows(self, flows: typing.Sequence[flow.Flow]) -> None: + """ + Replay server responses from flows. + """ + self.flowmap = {} for i in flows: - if i.response: + if i.response: # type: ignore l = self.flowmap.setdefault(self._hash(i), []) l.append(i) - - def clear(self): + ctx.master.addons.trigger("update", []) + + @command.command("replay.server.file") + def load_file(self, path: str) -> None: + try: + flows = io.read_flows_from_paths([path]) + except exceptions.FlowReadException as e: + raise exceptions.CommandError(str(e)) + self.load_flows(flows) + + @command.command("replay.server.stop") + def clear(self) -> None: + """ + Stop server replay. + """ self.flowmap = {} + ctx.master.addons.trigger("update", []) def count(self): return sum([len(i) for i in self.flowmap.values()]) @@ -38,27 +59,27 @@ class ServerPlayback: queriesArray = urllib.parse.parse_qsl(query, keep_blank_values=True) key = [str(r.port), str(r.scheme), str(r.method), str(path)] # type: List[Any] - if not self.options.server_replay_ignore_content: - if self.options.server_replay_ignore_payload_params and r.multipart_form: + if not ctx.options.server_replay_ignore_content: + if ctx.options.server_replay_ignore_payload_params and r.multipart_form: key.extend( (k, v) for k, v in r.multipart_form.items(multi=True) - if k.decode(errors="replace") not in self.options.server_replay_ignore_payload_params + if k.decode(errors="replace") not in ctx.options.server_replay_ignore_payload_params ) - elif self.options.server_replay_ignore_payload_params and r.urlencoded_form: + elif ctx.options.server_replay_ignore_payload_params and r.urlencoded_form: key.extend( (k, v) for k, v in r.urlencoded_form.items(multi=True) - if k not in self.options.server_replay_ignore_payload_params + if k not in ctx.options.server_replay_ignore_payload_params ) else: key.append(str(r.raw_content)) - if not self.options.server_replay_ignore_host: + if not ctx.options.server_replay_ignore_host: key.append(r.host) filtered = [] - ignore_params = self.options.server_replay_ignore_params or [] + ignore_params = ctx.options.server_replay_ignore_params or [] for p in queriesArray: if p[0] not in ignore_params: filtered.append(p) @@ -66,9 +87,9 @@ class ServerPlayback: key.append(p[0]) key.append(p[1]) - if self.options.server_replay_use_headers: + if ctx.options.server_replay_use_headers: headers = [] - for i in self.options.server_replay_use_headers: + for i in ctx.options.server_replay_use_headers: v = r.headers.get(i) headers.append((i, v)) key.append(headers) @@ -83,7 +104,7 @@ class ServerPlayback: """ hsh = self._hash(request) if hsh in self.flowmap: - if self.options.server_replay_nopop: + if ctx.options.server_replay_nopop: return self.flowmap[hsh][0] else: ret = self.flowmap[hsh].pop(0) @@ -91,16 +112,14 @@ class ServerPlayback: 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 configure(self, updated): + if not self.configured and ctx.options.server_replay: + self.configured = True + try: + flows = io.read_flows_from_paths(ctx.options.server_replay) + except exceptions.FlowReadException as e: + raise exceptions.OptionsError(str(e)) + self.load_flows(flows) def tick(self): if self.stop and not self.final_flow.live: @@ -112,13 +131,13 @@ class ServerPlayback: if rflow: response = rflow.response.copy() response.is_replay = True - if self.options.refresh_server_playback: + if ctx.options.refresh_server_playback: response.refresh() f.response = response if not self.flowmap: self.final_flow = f self.stop = True - elif self.options.replay_kill_extra: + elif ctx.options.replay_kill_extra: ctx.log.warn( "server_playback: killed non-replay request {}".format( f.request.url |