update client

client.py

@@ -4,6 +4,10 @@ import logging
 import sys
 import socket
 import re
+import time
+import os
+
+from client.ResponseHandler import ResponseHandler
 
 FORMAT = 'utf-8'
 BUFSIZE = 4096
@@ -50,6 +54,19 @@ def receive_bytes(client: socket.socket):
         yield buffer
 
 
+def receive(client: socket.socket):
+    if client.fileno() == -1:
+        raise Exception("Connection closed")
+
+    result = client.recv(BUFSIZE)
+
+    if len(result) == 0:
+        time.sleep(0.1)
+        result = client.recv(BUFSIZE)
+
+    return result
+
+
 def parse_header(data: bytes):
     headers = {}
 
@@ -108,6 +125,19 @@ def get_chunk(buffer: bytes):
     return buffer[:split_start], buffer[split_end:]
 
 
+def get_html_filename(headers):
+    if "CONTENT-LOCATION" not in headers:
+        return "index.html"
+
+    filename = headers["CONTENT-LOCATION"]
+    result = os.path.basename(filename).strip()
+
+    if len(result.strip()) == 0:
+        return 'index.html'
+
+    return result
+
+
 def response_parser(client: socket.socket):
     client.settimeout(3.0)
 
@@ -124,6 +154,33 @@ def response_parser(client: socket.socket):
         raise Exception("Invalid status-line")
     logging.debug("valid status-line: %r", status_line)
 
+    encoding = "plain"
+
+    if "TRANSFER-ENCODING" in headers:
+        encoding = headers["TRANSFER-ENCODING"]
+
+    if encoding == "plain" and "CONTENT-LENGTH" in headers:
+        payload_size = int(headers["CONTENT-LENGTH"])
+
+        if payload_size == 0:
+            return
+
+        filename = get_html_filename(headers)
+
+        f = open(filename, "wb")
+        f.write(buffer)
+
+        cur_payload_size = len(buffer)
+        while cur_payload_size < payload_size:
+            buffer = receive(client)
+            logging.debug("Received payload: %r", buffer)
+            if len(buffer) == 0:
+                logging.warning("Received payload length %s less than expected %s", payload_size, cur_payload_size)
+                break
+            cur_payload_size += len(buffer)
+            f.write(buffer)
+
+        f.close()
 
 
 def http_parser(client: socket.socket):
@@ -161,7 +218,7 @@ def main():
     client.sendall(message)
 
     response_parser(client)
-    http_parser(client)
+    # http_parser(client)
     # tmp = b''
     # keep = False
     # count = 0