diff --git a/fileserver.py b/fileserver.py index 23f3b74..f634334 100755 --- a/fileserver.py +++ b/fileserver.py @@ -36,11 +36,11 @@ class FileServerRequestHandler(BaseHTTPRequestHandler): # target += "?" if "?" not in target else "&" # target += queryStr - contentLength = self.headers.get('Content-Length') + content_length = self.headers.get('Content-Length') data = None - if contentLength and int(contentLength) > 0: - data = self.rfile.read(int(contentLength)) + if content_length and int(content_length) > 0: + data = self.rfile.read(int(content_length)) if "Host" in self.headers: del self.headers["Host"] @@ -114,11 +114,11 @@ class FileServerRequestHandler(BaseHTTPRequestHandler): self.wfile.write(data) if (path in self.server.dumpRequests or "/" in self.server.dumpRequests) and path != "/dummy": - contentLength = self.headers.get('Content-Length') + content_length = self.headers.get('Content-Length') body = None - if contentLength and int(contentLength) > 0: - body = self.rfile.read(int(contentLength)) + if content_length and int(content_length) > 0: + body = self.rfile.read(int(content_length)) print("===== Connection from:",self.client_address[0]) print("%s %s %s" % (self.command, self.path, self.request_version)) @@ -168,8 +168,9 @@ class HttpFileServer(HTTPServer): headers = { "Access-Control-Allow-Origin": "*", } + if mimeType: - headers["Content-Type"] = headers + headers["Content-Type"] = mimeType # return 200 - OK and data self.addRoute(name, lambda req: (200, data, headers)) diff --git a/util.py b/util.py index ac63bbb..081de98 100755 --- a/util.py +++ b/util.py @@ -107,13 +107,21 @@ def assert_not_empty(res, err=None): exit_with_error(res, err) def assert_content_contains(res, data, err=None): - util.assert_not_empty(res) + assert_not_empty(res) if isinstance(data, str) and data in res.text: return True elif data in res.content: return True - err = f"[-] '{res.url}' did not include '{data} in response" if err is None else err + err = f"[-] '{res.url}' did not include '{data}' in response" if err is None else err + exit_with_error(res, err) + +def assert_cookie_present(res, cookie_name, err=None): + assert_header_present(res, "Set-Cookie") + if cookie_name in res.cookies: + return True + + err = f"[-] '{res.url}' did not set-cookie '{cookie_name}'" if err is None else err exit_with_error(res, err) def assert_json_path(res, path, value, err=None):