Roman Hergenreder 7 kuukautta sitten
vanhempi
commit
29a3fd08d6
2 muutettua tiedostoa jossa 18 lisäystä ja 9 poistoa
  1. 8 7
      fileserver.py
  2. 10 2
      util.py

+ 8 - 7
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))

+ 10 - 2
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):