From 61ae5317a0185db363ad99da82821b6eb25b593d Mon Sep 17 00:00:00 2001 From: Alexis Delhaie Date: Mon, 5 Oct 2020 19:24:08 +0200 Subject: [PATCH] add headers --- src/main.py | 50 ++++++++++++++++++++++++++++++----------------- src/parameters.py | 17 ++++++++++++++++ src/thread.py | 19 +++++++++--------- 3 files changed, 59 insertions(+), 27 deletions(-) create mode 100644 src/parameters.py diff --git a/src/main.py b/src/main.py index b557c52..f6dc69c 100644 --- a/src/main.py +++ b/src/main.py @@ -8,11 +8,12 @@ from rich import box from rich.console import Console from rich.table import Table from thread import StressThread +from parameters import Options import thread import basic_producer_consumer console = Console(highlight=False) -VERSION = "1.3.2" +VERSION = "1.4" thread.console = console def main(): @@ -24,33 +25,35 @@ def main(): show_version() return - host = get_args("-h", True, 1001, 1002) - allow_ssl = get_flag("--ssl") - port = int(get_args("--port", False, 1003, 1004, "80")) - path = get_args("-p", True, 1005, 1006) - thread_number = int(get_args("-t", False, 1007, 1008, "5")) - timeout = int(get_args("-tm", False, 1009, 1010, "10")) - one_by_one = get_flag("--one-by-one") - ignore_available_threads = get_flag("--ignore-available-threads") - self_signed = get_flag("--allow-self-signed") + options = Options() + options.host = get_args("-h", True, 1001, 1002) + options.allow_ssl = get_flag("--ssl") + options.port = int(get_args("--port", False, 1003, 1004, "80")) + options.path = get_args("-p", False, 1005, 1006, "/") + options.thread_number = int(get_args("-t", False, 1007, 1008, "5")) + options.timeout = int(get_args("-tm", False, 1009, 1010, "10")) + options.one_by_one = get_flag("--one-by-one") + options.ignore_available_threads = get_flag("--ignore-available-threads") + options.self_signed = get_flag("--allow-self-signed") + options.headers = get_headers() - if one_by_one and ignore_available_threads: + if options.one_by_one and options.ignore_available_threads: console.print("[red]Error[/]: ambigous arguments, --one-by-one and --ignore-available-threads cannot be in the same command", style="bold") return - start(host, port, path, timeout, thread_number, allow_ssl, self_signed, one_by_one, ignore_available_threads) + start(options) -def start(host, port, path, timeout, thread_number, allow_ssl, self_signed, one_by_one = False, ignore_available_threads = False): +def start(options): thread_array = [] - for i in range(0, thread_number): - thread_array.append(StressThread(host, port, path, timeout, i, allow_ssl, self_signed, VERSION)) - if one_by_one: + for i in range(0, options.thread_number): + thread_array.append(StressThread(options, i, VERSION)) + if options.one_by_one: start_one_by_one(thread_array) - elif ignore_available_threads: + elif options.ignore_available_threads: start_all(thread_array) else: basic_producer_consumer.start(thread_array) - show_stat(thread_array, (timeout * 1000)) + show_stat(thread_array, (options.timeout * 1000)) # Run requests in one thread def start_one_by_one(threads): @@ -113,6 +116,7 @@ def show_help(): console.print(" --ignore-available-threads Ignore physical number of threads and start all request in the same time") console.print(" --ssl Use HTTPS/SSL") console.print(" --allow-self-signed Allow self signed SSL certificate") + console.print(" --header key=value Send a custom header (To add several headers, add several times the argument --header)") console.print(" --help") console.print(" /? Show this page") console.print(" --version Get information about the application and the system") @@ -150,6 +154,16 @@ def get_args(header, important, notfoundcode, valuenotfoundcode, default = ""): sys.exit(notfoundcode) return default +def get_headers(): + headers = {} + indices = [i for i, x in enumerate(sys.argv) if x == "--header"] + for i in indices: + if (len(sys.argv) > (i + 1)): + h_raw = sys.argv[i + 1] + key, value = h_raw.split('=', 1) + headers[key] = value + return headers + def get_flag(header): return (header in sys.argv) diff --git a/src/parameters.py b/src/parameters.py new file mode 100644 index 0000000..857f367 --- /dev/null +++ b/src/parameters.py @@ -0,0 +1,17 @@ + + +class Options: + + def __init__(self): + self.host = "" + self.allow_ssl = False + self.port = 80 + self.path = "/" + self.thread_number = 5 + self.timeout = 10 + self.one_by_one = False + self.ignore_available_threads = False + self.self_signed = False + self.headers = {} + + \ No newline at end of file diff --git a/src/thread.py b/src/thread.py index 23e4095..1df28d8 100644 --- a/src/thread.py +++ b/src/thread.py @@ -11,16 +11,17 @@ console = None class StressThread (threading.Thread): - def __init__(self, host, port, path, timeout, n, allow_ssl, self_signed, version): + def __init__(self, options, n, version): threading.Thread.__init__(self) self.user_agent = "PyStressTest/{}({} {} {})".format(version, platform.system(), os.name, platform.release()) - self.host = host - self.port = port - self.path = path - self.timeout = timeout + self.host = options.host + self.port = options.port + self.path = options.path + self.timeout = options.timeout self.n = n - self.allow_ssl = allow_ssl - self.self_signed = self_signed + self.allow_ssl = options.allow_ssl + self.self_signed = options.self_signed + self.headers = options.headers self.success = False self.time = 0 @@ -35,8 +36,8 @@ class StressThread (threading.Thread): c = http.client.HTTPSConnection(self.host, self.port, timeout=self.timeout, key_file=None, cert_file=None) else: c = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout) - headers = {"User-Agent": self.user_agent} - c.request(method="GET", url=self.path, headers=headers) + self.headers["User-Agent"] = self.user_agent + c.request(method="GET", url=self.path, headers=self.headers) res = c.getresponse() processed = round((time.time() - now) * 1000) self.print_result(res.status, res.reason, processed)