From 1966a174bb201331c5a52b214ba635fe350ea908 Mon Sep 17 00:00:00 2001
From: Arthur Bols <arthur@bols.dev>
Date: Thu, 18 Mar 2021 01:31:45 +0100
Subject: [PATCH] update client

---
 client.py                 | 59 ++++++++++++++++++++++++++++++++++++++-
 client/ResponseHandler.py | 42 ++++++++++++++++++++++++++++
 client/__init__.py        |  0
 3 files changed, 100 insertions(+), 1 deletion(-)
 create mode 100644 client/ResponseHandler.py
 create mode 100644 client/__init__.py

diff --git a/client.py b/client.py
index 2b188bf..3cd9cf1 100644
--- a/client.py
+++ b/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
diff --git a/client/ResponseHandler.py b/client/ResponseHandler.py
new file mode 100644
index 0000000..f690a09
--- /dev/null
+++ b/client/ResponseHandler.py
@@ -0,0 +1,42 @@
+import os
+from socket import socket
+from typing import Dict
+from urllib.parse import urlparse
+
+
+class ResponseHandler:
+    client: socket
+    url: str
+    headers: Dict[str, str]
+
+    def __init__(self, url: str, client: socket):
+        self.headers = {}
+        self.url = url
+        self.client = client
+        pass
+
+    def get_html_filename(self):
+        filename = "index.html"
+
+        parsed = urlparse(self.url)
+        if parsed.netloc == "":
+            parsed = urlparse("//" + self.url)
+        if len(parsed.path) != 0:
+            index = parsed.path.rfind("/")
+            if index == -1:
+                filename = parsed.path
+            elif parsed.path[-1] != "/":
+                filename = parsed.path[index:]
+
+        result = os.path.basename(filename).strip()
+        return result
+
+
+class PlainResponseHandler(ResponseHandler):
+    def __init__(self, url: str, client: socket):
+        super().__init__(url, client)
+
+
+class ChunkedResponseHandler(ResponseHandler):
+    def __init__(self, url: str, client: socket):
+        super().__init__(url, client)
diff --git a/client/__init__.py b/client/__init__.py
new file mode 100644
index 0000000..e69de29