This commit is contained in:
Maximilian Hils 2019-11-15 21:00:47 +01:00
parent a5412ab136
commit 8f2cee7225
2 changed files with 50 additions and 33 deletions

View File

@ -1,23 +1,23 @@
import queue import queue
import threading import threading
import typing
import time import time
import typing
from mitmproxy import log import mitmproxy.types
from mitmproxy import controller from mitmproxy import command
from mitmproxy import exceptions
from mitmproxy import http
from mitmproxy import flow
from mitmproxy import options
from mitmproxy import connections from mitmproxy import connections
from mitmproxy import controller
from mitmproxy import ctx
from mitmproxy import exceptions
from mitmproxy import flow
from mitmproxy import http
from mitmproxy import io
from mitmproxy import log
from mitmproxy import options
from mitmproxy.coretypes import basethread
from mitmproxy.net import server_spec, tls from mitmproxy.net import server_spec, tls
from mitmproxy.net.http import http1 from mitmproxy.net.http import http1
from mitmproxy.coretypes import basethread
from mitmproxy.utils import human from mitmproxy.utils import human
from mitmproxy import ctx
from mitmproxy import io
from mitmproxy import command
import mitmproxy.types
class RequestReplayThread(basethread.BaseThread): class RequestReplayThread(basethread.BaseThread):
@ -117,7 +117,7 @@ class RequestReplayThread(basethread.BaseThread):
finally: finally:
r.first_line_format = first_line_format_backup r.first_line_format = first_line_format_backup
f.live = False f.live = False
if server.connected(): if server and server.connected():
server.finish() server.finish()
server.close() server.close()

View File

@ -1,16 +1,19 @@
import hashlib import hashlib
import urllib
import typing import typing
import urllib
from mitmproxy import ctx
from mitmproxy import flow
from mitmproxy import exceptions
from mitmproxy import io
from mitmproxy import command
import mitmproxy.types import mitmproxy.types
from mitmproxy import command
from mitmproxy import ctx, http
from mitmproxy import exceptions
from mitmproxy import flow
from mitmproxy import io
class ServerPlayback: class ServerPlayback:
flowmap: typing.Dict[typing.Hashable, typing.List[http.HTTPFlow]]
configured: bool
def __init__(self): def __init__(self):
self.flowmap = {} self.flowmap = {}
self.configured = False self.configured = False
@ -82,10 +85,10 @@ class ServerPlayback:
Replay server responses from flows. Replay server responses from flows.
""" """
self.flowmap = {} self.flowmap = {}
for i in flows: for f in flows:
if i.response: # type: ignore if isinstance(f, http.HTTPFlow):
l = self.flowmap.setdefault(self._hash(i), []) lst = self.flowmap.setdefault(self._hash(f), [])
l.append(i) lst.append(f)
ctx.master.addons.trigger("update", []) ctx.master.addons.trigger("update", [])
@command.command("replay.server.file") @command.command("replay.server.file")
@ -108,12 +111,11 @@ class ServerPlayback:
def count(self) -> int: def count(self) -> int:
return sum([len(i) for i in self.flowmap.values()]) 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. Calculates a loose hash of the flow request.
""" """
r = flow.request r = flow.request
_, _, path, _, query, _ = urllib.parse.urlparse(r.url) _, _, path, _, query, _ = urllib.parse.urlparse(r.url)
queriesArray = urllib.parse.parse_qsl(query, keep_blank_values=True) queriesArray = urllib.parse.parse_qsl(query, keep_blank_values=True)
@ -158,20 +160,33 @@ class ServerPlayback:
repr(key).encode("utf8", "surrogateescape") repr(key).encode("utf8", "surrogateescape")
).digest() ).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 Returns the next flow object, or None if no matching flow was
found. found.
""" """
hsh = self._hash(request) request = flow.request
if hsh in self.flowmap: hash = self._hash(flow)
if hash in self.flowmap:
if ctx.options.server_replay_nopop: 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: else:
ret = self.flowmap[hsh].pop(0) ret = self.flowmap[hash].pop(0)
if not self.flowmap[hsh]: while not ret.response:
del self.flowmap[hsh] 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 return ret
else:
return None
def configure(self, updated): def configure(self, updated):
if not self.configured and ctx.options.server_replay: if not self.configured and ctx.options.server_replay:
@ -182,10 +197,11 @@ class ServerPlayback:
raise exceptions.OptionsError(str(e)) raise exceptions.OptionsError(str(e))
self.load_flows(flows) self.load_flows(flows)
def request(self, f): def request(self, f: http.HTTPFlow) -> None:
if self.flowmap: if self.flowmap:
rflow = self.next_flow(f) rflow = self.next_flow(f)
if rflow: if rflow:
assert rflow.response
response = rflow.response.copy() response = rflow.response.copy()
response.is_replay = True response.is_replay = True
if ctx.options.server_replay_refresh: if ctx.options.server_replay_refresh:
@ -197,4 +213,5 @@ class ServerPlayback:
f.request.url f.request.url
) )
) )
assert f.reply
f.reply.kill() f.reply.kill()