Add HTTP body size limit specification to command-line tools.

This commit is contained in:
Aldo Cortesi 2011-09-09 15:27:31 +12:00
parent 28daa93268
commit 67f2610032
8 changed files with 69 additions and 27 deletions

View File

@ -54,6 +54,12 @@ def common_options(parser):
action="store_true", dest="eventlog",
help="Show event log."
)
parser.add_option(
"-l",
action="store", dest="body_size_limit", default=None,
help="Byte size limit of HTTP request and response bodies."\
" Understands k/m/g suffixes, i.e. 3m for 3 megabytes."
)
parser.add_option(
"-n",
action="store_true", dest="no_server",

View File

@ -855,7 +855,6 @@ class Options(object):
"keepserving",
"kill",
"intercept",
"limit",
"no_server",
"refresh_server_playback",
"rfile",
@ -943,11 +942,6 @@ class ConsoleMaster(flow.FlowMaster):
self.conn_list_view = None
self.set_palette()
r = self.set_limit(options.limit)
if r:
print >> sys.stderr, "Limit error:", r
sys.exit(1)
r = self.set_intercept(options.intercept)
if r:
print >> sys.stderr, "Intercept error:", r

View File

@ -1260,7 +1260,7 @@ class FlowMaster(controller.Master):
f.response = None
f.error = None
self.process_new_request(f)
rt = proxy.RequestReplayThread(f, self.masterq)
rt = proxy.RequestReplayThread(f, self.masterq, None)
rt.start()
#end nocover

View File

@ -22,12 +22,13 @@ class ProxyError(Exception):
class ProxyConfig:
def __init__(self, certfile = None, ciphers = None, cacert = None, cert_wait_time=0):
def __init__(self, certfile = None, ciphers = None, cacert = None, cert_wait_time=0, body_size_limit = None):
self.certfile = certfile
self.ciphers = ciphers
self.cacert = cacert
self.certdir = None
self.cert_wait_time = cert_wait_time
self.body_size_limit = body_size_limit
def read_chunked(fp, limit):
@ -160,13 +161,13 @@ class FileLike:
#begin nocover
class RequestReplayThread(threading.Thread):
def __init__(self, flow, masterq):
self.flow, self.masterq = flow, masterq
def __init__(self, flow, masterq, body_size_limit):
self.flow, self.masterq, self.body_size_limit = flow, masterq, body_size_limit
threading.Thread.__init__(self)
def run(self):
try:
server = ServerConnection(self.flow.request)
server = ServerConnection(self.flow.request, self.body_size_limit)
server.send_request(self.flow.request)
response = server.read_response()
response._send(self.masterq)
@ -176,7 +177,8 @@ class RequestReplayThread(threading.Thread):
class ServerConnection:
def __init__(self, request):
def __init__(self, request, body_size_limit):
self.body_size_limit = body_size_limit
self.host = request.host
self.port = request.port
self.scheme = request.scheme
@ -226,7 +228,7 @@ class ServerConnection:
if self.request.method == "HEAD" or code == 204 or code == 304:
content = ""
else:
content = read_http_body(self.rfile, self, headers, True, None)
content = read_http_body(self.rfile, self, headers, True, self.body_size_limit)
return flow.Response(self.request, code, msg, headers, content)
def terminate(self):
@ -274,7 +276,7 @@ class ProxyHandler(SocketServer.StreamRequestHandler):
request = False
response = response._send(self.mqueue)
else:
server = ServerConnection(request)
server = ServerConnection(request, self.config.body_size_limit)
server.send_request(request)
try:
response = server.read_response()
@ -381,7 +383,7 @@ class ProxyHandler(SocketServer.StreamRequestHandler):
client_conn.close = True
if value == "keep-alive":
client_conn.close = False
content = read_http_body(self.rfile, client_conn, headers, False, None)
content = read_http_body(self.rfile, client_conn, headers, False, self.config.body_size_limit)
return flow.Request(client_conn, host, port, scheme, method, path, headers, content)
def send_response(self, response):
@ -463,7 +465,7 @@ def certificate_option_group(parser):
parser.add_option_group(group)
def process_certificate_option_group(parser, options):
def process_proxy_options(parser, options):
if options.cert:
options.cert = os.path.expanduser(options.cert)
if not os.path.exists(options.cert):
@ -475,9 +477,13 @@ def process_certificate_option_group(parser, options):
utils.dummy_ca(cacert)
if getattr(options, "cache", None) is not None:
options.cache = os.path.expanduser(options.cache)
body_size_limit = utils.parse_size(options.body_size_limit)
return ProxyConfig(
certfile = options.cert,
cacert = cacert,
ciphers = options.ciphers,
cert_wait_time = options.cert_wait_time
cert_wait_time = options.cert_wait_time,
body_size_limit = body_size_limit
)

View File

@ -410,3 +410,33 @@ def parse_url(url):
return scheme, host, port, path
def parse_size(s):
"""
Parses a size specification. Valid specifications are:
123: bytes
123k: kilobytes
123m: megabytes
123g: gigabytes
"""
if not s:
return None
mult = None
if s[-1].lower() == "k":
mult = 1024**1
elif s[-1].lower() == "m":
mult = 1024**2
elif s[-1].lower() == "g":
mult = 1024**3
if mult:
s = s[:-1]
else:
mult = 1
try:
return int(s) * mult
except ValueError:
raise ValueError("Invalid size specification: %s"%s)

View File

@ -38,12 +38,12 @@ if __name__ == '__main__':
if options.quiet:
options.verbose = 0
config = proxy.process_certificate_option_group(parser, options)
proxyconfig = proxy.process_proxy_options(parser, options)
if options.no_server:
server = None
else:
try:
server = proxy.ProxyServer(config, options.port, options.addr)
server = proxy.ProxyServer(proxyconfig, options.port, options.addr)
except proxy.ProxyServerError, v:
print >> sys.stderr, "mitmdump:", v.args[0]
sys.exit(1)

View File

@ -35,11 +35,6 @@ if __name__ == '__main__':
"Filters",
"See help in mitmproxy for filter expression syntax."
)
group.add_option(
"-l", "--limit", action="store",
type = "str", dest="limit", default=None,
help = "Limit filter expression."
)
group.add_option(
"-i", "--intercept", action="store",
type = "str", dest="intercept", default=None,
@ -48,7 +43,7 @@ if __name__ == '__main__':
parser.add_option_group(group)
options, args = parser.parse_args()
config = proxy.process_certificate_option_group(parser, options)
config = proxy.process_proxy_options(parser, options)
if options.no_server:
server = None
@ -61,7 +56,6 @@ if __name__ == '__main__':
opts = console.Options(**cmdline.get_common_options(options))
opts.intercept = options.intercept
opts.limit = options.limit
opts.debug = options.debug
m = console.ConsoleMaster(server, opts)

View File

@ -219,6 +219,17 @@ class u_parse_url(libpry.AutoTree):
assert not utils.parse_url("https://foo:bar")
assert not utils.parse_url("https://foo:")
class u_parse_size(libpry.AutoTree):
def test_simple(self):
assert utils.parse_size("1") == 1
assert utils.parse_size("1k") == 1024
assert utils.parse_size("1m") == 1024**2
assert utils.parse_size("1g") == 1024**3
libpry.raises(ValueError, utils.parse_size, "1f")
libpry.raises(ValueError, utils.parse_size, "ak")
tests = [
uformat_timestamp(),
uisBin(),
@ -233,5 +244,6 @@ tests = [
udummy_ca(),
udummy_cert(),
uLRUCache(),
u_parse_url()
u_parse_url(),
u_parse_size()
]