mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-27 02:24:18 +00:00
Refactor to extract ready_actions and write_values.
This commit is contained in:
Aldo Cortesi
parent
05f5e772c3
commit
d4ad3f0b2c
@ -33,7 +33,7 @@ class PathodHandler(tcp.BaseHandler):
|
||||
if not crafted and path.startswith(self.server.prefix):
|
||||
spec = urllib.unquote(path)[len(self.server.prefix):]
|
||||
try:
|
||||
crafted = rparse.parse(self.server.request_settings, spec)
|
||||
crafted = rparse.parse_response(self.server.request_settings, spec)
|
||||
except rparse.ParseException, v:
|
||||
crafted = rparse.InternalResponse(
|
||||
800,
|
||||
@ -95,7 +95,7 @@ class Pathod(tcp.TCPServer):
|
||||
except re.error:
|
||||
raise PathodError("Invalid regex in anchor: %s"%i[0])
|
||||
try:
|
||||
aresp = rparse.parse(self.request_settings, i[1])
|
||||
aresp = rparse.parse_response(self.request_settings, i[1])
|
||||
except rparse.ParseException, v:
|
||||
raise PathodError("Invalid page spec in anchor: '%s', %s"%(i[1], str(v)))
|
||||
self.anchors.append((arex, aresp))
|
||||
|
||||
@ -2,6 +2,8 @@ import operator, string, random, mmap, os, time
|
||||
import contrib.pyparsing as pp
|
||||
from netlib import http_status
|
||||
|
||||
BLOCKSIZE = 1024
|
||||
|
||||
class ParseException(Exception):
|
||||
def __init__(self, msg, s, col):
|
||||
Exception.__init__(self)
|
||||
@ -19,6 +21,52 @@ class ParseException(Exception):
|
||||
class ServerError(Exception): pass
|
||||
|
||||
|
||||
def ready_actions(l, lst):
|
||||
ret = []
|
||||
for i in lst:
|
||||
itms = list(i)
|
||||
if i[0] == "r":
|
||||
itms[0] = random.randrange(l)
|
||||
if i[0] == "a":
|
||||
itms[0] = l+1
|
||||
ret.append(tuple(itms))
|
||||
ret.sort()
|
||||
return ret
|
||||
|
||||
|
||||
def write_values(fp, vals, actions, sofar=0, skip=0, blocksize=BLOCKSIZE):
|
||||
"""
|
||||
vals: A list of values, which may be strings or Value objects.
|
||||
actions: A list of (offset, action, arg) tuples. Action may be "pause" or "disconnect".
|
||||
|
||||
Return True if connection should disconnect.
|
||||
"""
|
||||
while vals:
|
||||
part = vals.pop()
|
||||
for i in range(skip, len(part), blocksize):
|
||||
d = part[i:i+blocksize]
|
||||
if actions and actions[-1][0] < (sofar + len(d)):
|
||||
p = actions.pop()
|
||||
offset = p[0]-sofar
|
||||
vals.append(part)
|
||||
if p[1] == "pause":
|
||||
fp.write(d[:offset])
|
||||
time.sleep(p[2])
|
||||
return write_values(
|
||||
fp, vals, actions,
|
||||
sofar=sofar+offset,
|
||||
skip=i+offset,
|
||||
blocksize=blocksize
|
||||
)
|
||||
elif p[1] == "disconnect":
|
||||
fp.write(d[:offset])
|
||||
return True
|
||||
fp.write(d)
|
||||
sofar += len(d)
|
||||
skip = 0
|
||||
|
||||
|
||||
|
||||
DATATYPES = dict(
|
||||
ascii_letters = string.ascii_letters,
|
||||
ascii_lowercase = string.ascii_lowercase,
|
||||
@ -328,7 +376,6 @@ class Code:
|
||||
return e.setParseAction(lambda x: klass(*x))
|
||||
|
||||
|
||||
BLOCKSIZE = 1024
|
||||
class Response:
|
||||
comps = (
|
||||
Body,
|
||||
@ -375,46 +422,6 @@ class Response:
|
||||
l += len(self.body)
|
||||
return l
|
||||
|
||||
def ready_actions(self, l, lst):
|
||||
ret = []
|
||||
for i in lst:
|
||||
itms = list(i)
|
||||
if i[0] == "r":
|
||||
itms[0] = random.randrange(l)
|
||||
if i[0] == "a":
|
||||
itms[0] = l+1
|
||||
ret.append(tuple(itms))
|
||||
ret.sort()
|
||||
return ret
|
||||
|
||||
def write_values(self, fp, vals, actions, sofar=0, skip=0, blocksize=BLOCKSIZE):
|
||||
"""
|
||||
Return True if connection should disconnect.
|
||||
"""
|
||||
while vals:
|
||||
part = vals.pop()
|
||||
for i in range(skip, len(part), blocksize):
|
||||
d = part[i:i+blocksize]
|
||||
if actions and actions[-1][0] < (sofar + len(d)):
|
||||
p = actions.pop()
|
||||
offset = p[0]-sofar
|
||||
vals.append(part)
|
||||
if p[1] == "pause":
|
||||
fp.write(d[:offset])
|
||||
time.sleep(p[2])
|
||||
return self.write_values(
|
||||
fp, vals, actions,
|
||||
sofar=sofar+offset,
|
||||
skip=i+offset,
|
||||
blocksize=blocksize
|
||||
)
|
||||
elif p[1] == "disconnect":
|
||||
fp.write(d[:offset])
|
||||
return True
|
||||
fp.write(d)
|
||||
sofar += len(d)
|
||||
skip = 0
|
||||
|
||||
def serve(self, fp):
|
||||
started = time.time()
|
||||
if self.body and not self.get_header("Content-Length"):
|
||||
@ -443,9 +450,9 @@ class Response:
|
||||
if self.body:
|
||||
vals.append(self.body)
|
||||
vals.reverse()
|
||||
actions = self.ready_actions(self.length(), self.actions)
|
||||
actions = ready_actions(self.length(), self.actions)
|
||||
actions.reverse()
|
||||
disconnect = self.write_values(fp, vals, actions[:])
|
||||
disconnect = write_values(fp, vals, actions[:])
|
||||
duration = time.time() - started
|
||||
return dict(
|
||||
disconnect = disconnect,
|
||||
@ -498,7 +505,7 @@ class InternalResponse(Response):
|
||||
return d
|
||||
|
||||
|
||||
def parse(settings, s):
|
||||
def parse_response(settings, s):
|
||||
try:
|
||||
return CraftedResponse(settings, s, Response.expr().parseString(s, parseAll=True))
|
||||
except pp.ParseException, v:
|
||||
|
||||
@ -137,9 +137,9 @@ class TestMisc:
|
||||
|
||||
|
||||
class TestDisconnects:
|
||||
def test_parse(self):
|
||||
assert (0, "disconnect") in rparse.parse({}, "400:d0").actions
|
||||
assert ("r", "disconnect") in rparse.parse({}, "400:dr").actions
|
||||
def test_parse_response(self):
|
||||
assert (0, "disconnect") in rparse.parse_response({}, "400:d0").actions
|
||||
assert ("r", "disconnect") in rparse.parse_response({}, "400:dr").actions
|
||||
|
||||
def test_at(self):
|
||||
e = rparse.DisconnectAt.expr()
|
||||
@ -156,13 +156,13 @@ class TestDisconnects:
|
||||
|
||||
|
||||
class TestShortcuts:
|
||||
def test_parse(self):
|
||||
assert rparse.parse({}, "400:c'foo'").headers[0][0][:] == "Content-Type"
|
||||
assert rparse.parse({}, "400:l'foo'").headers[0][0][:] == "Location"
|
||||
def test_parse_response(self):
|
||||
assert rparse.parse_response({}, "400:c'foo'").headers[0][0][:] == "Content-Type"
|
||||
assert rparse.parse_response({}, "400:l'foo'").headers[0][0][:] == "Location"
|
||||
|
||||
|
||||
class TestPauses:
|
||||
def test_parse(self):
|
||||
def test_parse_response(self):
|
||||
e = rparse.PauseAt.expr()
|
||||
v = e.parseString("p10,10")[0]
|
||||
assert v.seconds == 10
|
||||
@ -178,109 +178,105 @@ class TestPauses:
|
||||
assert v.offset == "a"
|
||||
|
||||
def test_request(self):
|
||||
r = rparse.parse({}, '400:p10,10')
|
||||
r = rparse.parse_response({}, '400:p10,10')
|
||||
assert r.actions[0] == (10, "pause", 10)
|
||||
|
||||
|
||||
class TestParse:
|
||||
def test_parse_err(self):
|
||||
tutils.raises(rparse.ParseException, rparse.parse, {}, "400:msg,b:")
|
||||
tutils.raises(rparse.ParseException, rparse.parse_response, {}, "400:msg,b:")
|
||||
try:
|
||||
rparse.parse({}, "400'msg':b:")
|
||||
rparse.parse_response({}, "400'msg':b:")
|
||||
except rparse.ParseException, v:
|
||||
assert v.marked()
|
||||
assert str(v)
|
||||
|
||||
def test_parse_header(self):
|
||||
r = rparse.parse({}, '400:h"foo"="bar"')
|
||||
r = rparse.parse_response({}, '400:h"foo"="bar"')
|
||||
assert r.get_header("foo") == "bar"
|
||||
|
||||
def test_parse_pause_before(self):
|
||||
r = rparse.parse({}, "400:p10,0")
|
||||
r = rparse.parse_response({}, "400:p10,0")
|
||||
assert (0, "pause", 10) in r.actions
|
||||
|
||||
def test_parse_pause_after(self):
|
||||
r = rparse.parse({}, "400:p10,a")
|
||||
r = rparse.parse_response({}, "400:p10,a")
|
||||
assert ("a", "pause", 10) in r.actions
|
||||
|
||||
def test_parse_pause_random(self):
|
||||
r = rparse.parse({}, "400:p10,r")
|
||||
r = rparse.parse_response({}, "400:p10,r")
|
||||
assert ("r", "pause", 10) in r.actions
|
||||
|
||||
def test_parse_stress(self):
|
||||
r = rparse.parse({}, "400:b@100g")
|
||||
r = rparse.parse_response({}, "400:b@100g")
|
||||
assert r.length()
|
||||
|
||||
|
||||
class TestResponse:
|
||||
def dummy_response(self):
|
||||
return rparse.parse({}, "400'msg'")
|
||||
|
||||
def test_response(self):
|
||||
r = rparse.parse({}, "400'msg'")
|
||||
assert r.code == 400
|
||||
assert r.msg == "msg"
|
||||
|
||||
r = rparse.parse({}, "400'msg':b@100b")
|
||||
assert r.msg == "msg"
|
||||
assert r.body[:]
|
||||
assert str(r)
|
||||
|
||||
def test_ready_actions(self):
|
||||
r = rparse.parse({}, "400'msg'")
|
||||
|
||||
x = [(0, 5)]
|
||||
assert r.ready_actions(100, x) == x
|
||||
|
||||
x = [("r", 5)]
|
||||
ret = r.ready_actions(100, x)
|
||||
assert 0 <= ret[0][0] < 100
|
||||
|
||||
x = [("a", "pause", 5)]
|
||||
ret = r.ready_actions(100, x)
|
||||
assert ret[0][0] > 100
|
||||
|
||||
x = [(1, 5), (0, 5)]
|
||||
assert r.ready_actions(100, x) == sorted(x)
|
||||
|
||||
class TestWriteValues:
|
||||
def test_write_values_disconnects(self):
|
||||
r = self.dummy_response()
|
||||
s = cStringIO.StringIO()
|
||||
tst = "foo"*100
|
||||
r.write_values(s, [tst], [(0, "disconnect")], blocksize=5)
|
||||
rparse.write_values(s, [tst], [(0, "disconnect")], blocksize=5)
|
||||
assert not s.getvalue()
|
||||
|
||||
def test_write_values(self):
|
||||
tst = "foo"*1025
|
||||
r = rparse.parse({}, "400'msg'")
|
||||
|
||||
s = cStringIO.StringIO()
|
||||
r.write_values(s, [tst], [])
|
||||
rparse.write_values(s, [tst], [])
|
||||
assert s.getvalue() == tst
|
||||
|
||||
def test_write_values_pauses(self):
|
||||
tst = "".join(str(i) for i in range(10))
|
||||
r = rparse.parse({}, "400'msg'")
|
||||
|
||||
for i in range(2, 10):
|
||||
s = cStringIO.StringIO()
|
||||
r.write_values(s, [tst], [(2, "pause", 0), (1, "pause", 0)], blocksize=i)
|
||||
rparse.write_values(s, [tst], [(2, "pause", 0), (1, "pause", 0)], blocksize=i)
|
||||
assert s.getvalue() == tst
|
||||
|
||||
for i in range(2, 10):
|
||||
s = cStringIO.StringIO()
|
||||
r.write_values(s, [tst], [(1, "pause", 0)], blocksize=i)
|
||||
rparse.write_values(s, [tst], [(1, "pause", 0)], blocksize=i)
|
||||
assert s.getvalue() == tst
|
||||
|
||||
tst = ["".join(str(i) for i in range(10))]*5
|
||||
for i in range(2, 10):
|
||||
s = cStringIO.StringIO()
|
||||
r.write_values(s, tst[:], [(1, "pause", 0)], blocksize=i)
|
||||
rparse.write_values(s, tst[:], [(1, "pause", 0)], blocksize=i)
|
||||
assert s.getvalue() == "".join(tst)
|
||||
|
||||
|
||||
def test_ready_actions():
|
||||
x = [(0, 5)]
|
||||
assert rparse.ready_actions(100, x) == x
|
||||
|
||||
x = [("r", 5)]
|
||||
ret = rparse.ready_actions(100, x)
|
||||
assert 0 <= ret[0][0] < 100
|
||||
|
||||
x = [("a", "pause", 5)]
|
||||
ret = rparse.ready_actions(100, x)
|
||||
assert ret[0][0] > 100
|
||||
|
||||
x = [(1, 5), (0, 5)]
|
||||
assert rparse.ready_actions(100, x) == sorted(x)
|
||||
|
||||
|
||||
class TestResponse:
|
||||
def dummy_response(self):
|
||||
return rparse.parse_response({}, "400'msg'")
|
||||
|
||||
def test_response(self):
|
||||
r = rparse.parse_response({}, "400'msg'")
|
||||
assert r.code == 400
|
||||
assert r.msg == "msg"
|
||||
|
||||
r = rparse.parse_response({}, "400'msg':b@100b")
|
||||
assert r.msg == "msg"
|
||||
assert r.body[:]
|
||||
assert str(r)
|
||||
|
||||
def test_render(self):
|
||||
s = cStringIO.StringIO()
|
||||
r = rparse.parse({}, "400'msg'")
|
||||
r = rparse.parse_response({}, "400'msg'")
|
||||
assert r.serve(s)
|
||||
|
||||
def test_length(self):
|
||||
@ -288,6 +284,6 @@ class TestResponse:
|
||||
s = cStringIO.StringIO()
|
||||
x.serve(s)
|
||||
assert x.length() == len(s.getvalue())
|
||||
testlen(rparse.parse({}, "400'msg'"))
|
||||
testlen(rparse.parse({}, "400'msg':h'foo'='bar'"))
|
||||
testlen(rparse.parse({}, "400'msg':h'foo'='bar':b@100b"))
|
||||
testlen(rparse.parse_response({}, "400'msg'"))
|
||||
testlen(rparse.parse_response({}, "400'msg':h'foo'='bar'"))
|
||||
testlen(rparse.parse_response({}, "400'msg':h'foo'='bar':b@100b"))
|
||||
|
||||
Loading…
Reference in New Issue
Block a user