Update

client/command.py

@@ -1,16 +1,39 @@
 import logging
 from abc import ABC, abstractmethod
+from typing import Dict, Tuple
 from urllib.parse import urlparse
 
-from client.response_handler import ResponseHandler
 from client.httpclient import FORMAT, HTTPClient
 from httplib import parser
 from httplib.exceptions import InvalidResponse, InvalidStatusLine, UnsupportedEncoding
+from httplib.message import Message
+from httplib.retriever import PreambleRetriever
+
+sockets: Dict[str, HTTPClient] = {}
+
+
+def create(command: str, url: str, port):
+    if command == "GET":
+        return GetCommand(url, port)
+    elif command == "HEAD":
+        return HeadCommand(url, port)
+    elif command == "POST":
+        return PostCommand(url, port)
+    elif command == "PUT":
+        return PutCommand(url, port)
+    else:
+        raise ValueError()
+
 
 class AbstractCommand(ABC):
+    uri: str
+    host: str
+    path: str
+    port: Tuple[str, int]
 
-    def __init__(self, url: str, port: str):
-        self.url = url
+    def __init__(self, uri: str, port):
+        self.uri = uri
+        self.host, _, self.path = parser.parse_uri(uri)
         self.port = port
 
     @property
@@ -18,20 +41,6 @@ class AbstractCommand(ABC):
     def command(self):
         pass
 
-    @staticmethod
-    def create(command: str, url: str, port: str):
-        if command == "GET":
-            return GetCommand(url, port)
-        elif command == "HEAD":
-            return HeadCommand(url, port)
-        elif command == "POST":
-            return PostCommand(url, port)
-        elif command == "PUT":
-            return PutCommand(url, port)
-        else:
-            raise ValueError()
-
-
     @staticmethod
     def build_message(command, host, path):
         message = f"{command} {path} HTTP/1.1\r\n"
@@ -40,26 +49,34 @@ class AbstractCommand(ABC):
 
         return message.encode(FORMAT)
 
-    def execute(self):
+    def execute(self, sub_request=False):
         (host, path) = self.parse_uri()
 
-        client = HTTPClient(host)
-        client.conn.connect((host, int(self.port)))
+        client = sockets.get(host)
+
+        if client and client.is_closed():
+            sockets.pop(self.host)
+            client = None
+
+        if not client:
+            client = HTTPClient(host)
+            client.conn.connect((host, self.port))
+            sockets[host] = client
 
         message = f"{self.command} {path} HTTP/1.1\r\n"
-        message += f"Host: {host}\r\n"
+        message += f"Host: {host}:{self.port}\r\n"
         message += "Accept: */*\r\nAccept-Encoding: identity\r\n"
         encoded_msg = self._build_message(message)
 
-        logging.info("---request begin---\r\n%s---request end---", encoded_msg.decode(FORMAT))
+        logging.debug("---request begin---\r\n%s---request end---", encoded_msg.decode(FORMAT))
 
-        logging.debug("Sending HTTP message: %r", encoded_msg)
         client.conn.sendall(encoded_msg)
 
         logging.info("HTTP request sent, awaiting response...")
 
         try:
-            self._await_response(client)
+            retriever = PreambleRetriever(client)
+            self._await_response(client, retriever)
         except InvalidResponse as e:
             logging.debug("Internal error: Response could not be parsed", exc_info=e)
             return
@@ -69,9 +86,10 @@ class AbstractCommand(ABC):
         except UnsupportedEncoding as e:
             logging.debug("Internal error: Unsupported encoding in response", exc_info=e)
         finally:
-            client.close()
+            if not sub_request:
+                client.close()
 
-    def _await_response(self, client):
+    def _await_response(self, client, retriever):
         while True:
             line = client.read_line()
             print(line, end="")
@@ -82,11 +100,11 @@ class AbstractCommand(ABC):
         return (message + "\r\n").encode(FORMAT)
 
     def parse_uri(self):
-        parsed = urlparse(self.url)
+        parsed = urlparse(self.uri)
 
         # If there is no netloc, the url is invalid, so prepend `//` and try again
         if parsed.netloc == "":
-            parsed = urlparse("//" + self.url)
+            parsed = urlparse("//" + self.uri)
 
         host = parsed.netloc
         path = parsed.path
@@ -105,6 +123,7 @@ class AbstractWithBodyCommand(AbstractCommand, ABC):
     @staticmethod
     def build_message(command, host, path):
         message = AbstractCommand.build_message()
+
     def _build_message(self, message: str) -> bytes:
         body = input(f"Enter {self.command} data: ").encode(FORMAT)
         print()
@@ -126,18 +145,31 @@ class HeadCommand(AbstractCommand):
 
 
 class GetCommand(AbstractCommand):
+
+    def __init__(self, uri: str, port, dir=None):
+        super().__init__(uri, port)
+        self.dir = dir
+        self.filename = None
+
     @property
     def command(self):
         return "GET"
 
-    def _await_response(self, client):
-        (version, status, msg) = parser.get_status_line(client)
-        logging.debug("Parsed status-line: version: %s, status: %s", version, status)
-        headers = parser.get_headers(client)
-        logging.debug("Parsed headers: %r", headers)
+    def _get_preamble(self, retriever):
+        lines = retriever.retrieve()
+        (version, status, msg) = parser.parse_status_line(next(lines))
+        headers = parser.parse_headers(lines)
 
-        handler = ResponseHandler.create(client, headers, status, self.url)
-        handler.handle()
+        logging.debug("---response begin---\r\n%s--- response end---", "".join(retriever.buffer))
+
+        return Message(version, status, msg, headers)
+
+    def _await_response(self, client, retriever) -> str:
+        msg = self._get_preamble(retriever)
+
+        from client import response_handler
+        self.filename = response_handler.handle(client, msg, self, self.dir)
+        return
 
 
 class PostCommand(AbstractWithBodyCommand):