From 29a3fd08d6f325498ad6b1e37195bbb050b7acdf Mon Sep 17 00:00:00 2001
From: Roman Hergenreder <mail@romanh.de>
Date: Mon, 16 Oct 2023 17:44:22 +0200
Subject: [PATCH] bugfix

---
 fileserver.py | 15 ++++++++-------
 util.py       | 12 ++++++++++--
 2 files changed, 18 insertions(+), 9 deletions(-)

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):