contrib: Make the failures mitmproxy script a bit smarter.

master
Damir Jelić 2021-03-02 12:24:23 +01:00
parent 83926e154b
commit 56a696d1c0
1 changed files with 37 additions and 23 deletions

View File

@ -11,6 +11,7 @@ To run execute it with mitmproxy:
"""
import time
import json
import random
from mitmproxy import http
from mitmproxy.script import concurrent
@ -18,30 +19,43 @@ from mitmproxy.script import concurrent
REQUEST_COUNT = 0
@concurrent
def request(flow):
global REQUEST_COUNT
def timeout(flow):
timeout = 60 if "sync" in flow.request.pretty_url else 30
time.sleep(timeout)
return None
REQUEST_COUNT += 1
if REQUEST_COUNT % 2 == 0:
return
elif REQUEST_COUNT % 3 == 0:
flow.response = http.HTTPResponse.make(
500,
b"Gateway error",
)
elif REQUEST_COUNT % 7 == 0:
if "sync" in flow.request.pretty_url:
time.sleep(60)
else:
time.sleep(30)
else:
flow.response = http.HTTPResponse.make(
# A map holding our failure modes.
# The keys are just descriptive names for the failure mode while the values
# hold a tuple containing a function that may or may not create a failure and
# the probability weight at which rate this failure should be triggered.
#
# The method should return an http.HTTPResponse if it should modify the
# response or None if the response should be passed as is.
FAILURES = {
"Success": (lambda x: None, 50),
"Gateway error":
(lambda _: http.HTTPResponse.make(500, b"Gateway error"), 20),
"Limit exeeded": (lambda _: http.HTTPResponse.make(
429,
json.dumps({
"errcode": "M_LIMIT_EXCEEDED",
"error": "Too many requests",
"retry_after_ms": 2000
})
)
})), 20),
"Timeout error": (timeout, 10)
}
@concurrent
def request(flow):
global FAILURES
weights = [weight for (_, weight) in FAILURES.values()]
failure = random.choices(list(FAILURES), weights=weights)[0]
failure_func, _ = FAILURES[failure]
response = failure_func(flow)
if response:
flow.response = response