small update

This commit is contained in:
2021-03-28 00:35:58 +01:00
parent 5b5a821522
commit a3ce68330f
2 changed files with 27 additions and 19 deletions

View File

@@ -86,7 +86,7 @@ class BasicResponseHandler(ResponseHandler):
if 300 <= self.msg.status < 400: if 300 <= self.msg.status < 400:
# Redirect # Redirect
self._skip_body() self._skip_body()
return self._do_handle_redirect() return self._handle_redirect()
if 400 <= self.msg.status < 600: if 400 <= self.msg.status < 600:
self._skip_body() self._skip_body()
# Dump headers and exit with error # Dump headers and exit with error
@@ -96,7 +96,7 @@ class BasicResponseHandler(ResponseHandler):
return None return None
def _do_handle_redirect(self): def _handle_redirect(self):
if self.msg.status == 304: if self.msg.status == 304:
print("".join(self.msg.raw), end="") print("".join(self.msg.raw), end="")
return None return None
@@ -157,7 +157,8 @@ class DownloadHandler(ResponseHandler, ABC):
return tmp_path return tmp_path
def get_filename(self): def get_filename(self):
"""Returns the filename to download the payload to. """
Returns the filename to download the payload to.
""" """
filename = os.path.basename(self.cmd.path) filename = os.path.basename(self.cmd.path)
if filename == '': if filename == '':
@@ -211,6 +212,13 @@ class HTMLDownloadHandler(DownloadHandler):
return self.path return self.path
def _download_images(self, tmp_filename, target_filename, charset=FORMAT): def _download_images(self, tmp_filename, target_filename, charset=FORMAT):
"""
Downloads images referenced in the html of `tmp_filename` and replaces the references in the html
and writes it to `target_filename`.
@param tmp_filename: the path to the temporary html file
@param target_filename: the path for the final html fil
@param charset: the charset to decode `tmp_filename`
"""
try: try:
fp = open(tmp_filename, "r", encoding=charset) fp = open(tmp_filename, "r", encoding=charset)

View File

@@ -89,7 +89,7 @@ class AbstractCommand(ABC):
if content_type: if content_type:
message += f"Content-Type: {content_type}" message += f"Content-Type: {content_type}"
if content_type.startswith("text"): if content_type.startswith("text"):
message += "; charset=UTF-8" message += f"; charset={FORMAT}"
message += "\r\n" message += "\r\n"
elif content_length > 0: elif content_length > 0:
message += f"Content-Type: application/octet-stream\r\n" message += f"Content-Type: application/octet-stream\r\n"
@@ -141,6 +141,20 @@ class AbstractCommand(ABC):
return True return True
def get_mimetype(self, path):
mime = mimetypes.guess_type(path)[0]
if mime:
return mime
try:
file = open(path, "r", encoding=FORMAT)
file.readline()
file.close()
return "text/plain"
except UnicodeDecodeError:
return "application/octet-stream"
class AbstractModifyCommand(AbstractCommand, ABC): class AbstractModifyCommand(AbstractCommand, ABC):
@@ -191,7 +205,7 @@ class HeadCommand(AbstractCommand):
def execute(self): def execute(self):
path = self._get_path() path = self._get_path()
mime = mimetypes.guess_type(path)[0] mime = self.get_mimetype(path)
return self._build_message(200, mime, b"") return self._build_message(200, mime, b"")
@@ -204,20 +218,6 @@ class GetCommand(AbstractCommand):
def _conditional_headers(self): def _conditional_headers(self):
return {'if-modified-since': self._if_modified_since} return {'if-modified-since': self._if_modified_since}
def get_mimetype(self, path):
mime = mimetypes.guess_type(path)[0]
if mime:
return mime
try:
file = open(path, "r", encoding="utf-8")
file.readline()
file.close()
return "text/plain"
except UnicodeDecodeError:
return "application/octet-stream"
def execute(self): def execute(self):
path = self._get_path() path = self._get_path()
mime = self.get_mimetype(path) mime = self.get_mimetype(path)