From a27b25182b659a1ce4c1ffef2f22720ed8eeaa94 Mon Sep 17 00:00:00 2001
From: dece <shgck@pistache.land>
Date: Sun, 10 Jul 2022 18:24:59 +0200
Subject: [PATCH] shrlok: add html type

---
 .gitignore       |  1 +
 shrlok/shrlok.py | 15 +++++++++++----
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/.gitignore b/.gitignore
index 25aacff..45c0524 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+__pycache__/
 build/
 dist/
 *.egg-info/
diff --git a/shrlok/shrlok.py b/shrlok/shrlok.py
index 70ce367..9d991be 100644
--- a/shrlok/shrlok.py
+++ b/shrlok/shrlok.py
@@ -20,7 +20,8 @@ The content is 13 bytes (assuming no LF at the end), the header is 14 bytes,
 plus the null byte it is 13 + 14 + 1 = 28 bytes, thus the prefixed length.
 
 After the content is succesfully retrieved and put on an appropriate location,
-the server will reply the file path through the socket and close the connection.
+the server will reply the file path through the socket and close the
+connection.
 """
 
 import argparse
@@ -66,6 +67,8 @@ class Handler(socketserver.StreamRequestHandler):
         file_name = None
         if header.get("type") == "txt":
             file_name = write_text(data, title=header.get("title"))
+        elif header.get("type") == "html":
+            file_name = write_content(data)
         else:
             print("Unknown type.")
         if file_name is None:
@@ -99,18 +102,22 @@ class Handler(socketserver.StreamRequestHandler):
 
 def write_text(data: bytes, title=None):
     content = "<pre>{}</pre>".format(data.decode(errors="replace"))
-    html = HTML_TEMPLATE.format(title=title or "?", content=content)
+    html = HTML_TEMPLATE.format(title=title or "?", content=content).encode()
+    write_content(html)
+
+
+def write_content(data: bytes):
     path = str(Path(ARGS.root) / "txt")
     if not os.path.isdir(path):
         print(f"Missing '{path}' folder.")
         return None
     with tempfile.NamedTemporaryFile(
-            mode="wt",
+            mode="wb",
             dir=path,
             suffix=".html",
             delete=False
     ) as output_file:
-        output_file.write(html)
+        output_file.write(data)
         os.chmod(output_file.name, 0o644)
         return output_file.name