From b51b314d747ee599ac9236bd0ed33e927d7a822f Mon Sep 17 00:00:00 2001
From: dece <shgck@pistache.land>
Date: Sat, 3 Sep 2022 20:01:59 +0200
Subject: [PATCH] improoooove

---
 share.sh         | 23 +++++++++++++++++++----
 shrlok/shrlok.py | 34 ++++++++++++++++++++++++----------
 2 files changed, 43 insertions(+), 14 deletions(-)

diff --git a/share.sh b/share.sh
index a5d2ebe..f93b8f1 100755
--- a/share.sh
+++ b/share.sh
@@ -1,15 +1,20 @@
 #!/bin/bash
 # Share data to a running Shrlok instance.
 
+socket_path="/run/shrlok/shr.sock"
+
 usage() {
-    echo "Usage: $0 [-s SOCKET_PATH] [] TYPE FILE"
+    echo "Usage: $0 [-s SOCKET_PATH] [-o key1=value1,key2=value2,…] TYPE FILE"
+    echo "  -s  override shrlok socket path (default: $socket_path)"
+    echo "  -o  comma separated "k=v" pairs to add to the shrlok header"
 }
 
-socket_path="/run/shrlok/shr.sock"
-while getopts "hs:" OPTION; do
+header_args=
+while getopts "hs:o:" OPTION; do
     case $OPTION in
         h) usage; exit 0 ;;
         s) socket_path="$OPTARG" ;;
+        o) readarray -d "," -t header_args <<< "$OPTARG" ;;
         *) usage; exit 1 ;;
     esac
 done
@@ -24,10 +29,20 @@ if [ ! -S "$socket_path" ]; then
     exit 1
 fi
 
+header_args_str=""
+if [[ -n "$header_args" ]]; then
+    for item in "${header_args[@]}"; do
+        readarray -d "=" -t kv <<< "$item"
+        key="${kv[0]}"
+        val="$(echo "${kv[1]}" | xargs)"  # strip trailing LFs
+        header_args_str+=",\"$key\":\"$val\""
+    done
+fi
+
 # The first thing sent through the socket is the total packet length, as ASCII
 # digits for the convenience of scripts. Build the header before hand to get its
 # size, then get the file length and push everything to the socket with socat.
-header="{\"type\":\"$share_type\"}"
+header="{\"type\":\"$share_type\"$header_args_str}"
 header_length="${#header}"
 file_length="$(wc -c "$input_file" | cut -d' ' -f1)"
 full_length="$(( $header_length + 1 + $file_length ))"
diff --git a/shrlok/shrlok.py b/shrlok/shrlok.py
index f1ebd75..47412ab 100644
--- a/shrlok/shrlok.py
+++ b/shrlok/shrlok.py
@@ -66,9 +66,13 @@ class Handler(socketserver.StreamRequestHandler):
 
         file_name = None
         if header.get("type") == "txt":
-            file_name = write_text(data, title=header.get("title"))
+            file_name = write_text(data, title=header.get("title", ""))
         elif header.get("type") == "raw":
-            file_name = write_content(data)
+            file_name = write_content(
+                data,
+                name=header.get("name", ""),
+                extension=header.get("ext", ""),
+            )
         else:
             print("Unknown type.")
         if file_name is None:
@@ -100,26 +104,36 @@ class Handler(socketserver.StreamRequestHandler):
         return data
 
 
-def write_text(data: bytes, title=None):
+def write_text(data: bytes, title: str = "", name: str = ""):
     content = "<pre>{}</pre>".format(data.decode(errors="replace"))
     html = HTML_TEMPLATE.format(title=title or "", content=content).encode()
-    return write_content(html)
+    return write_content(html, extension="html")
 
 
-def write_content(data: bytes):
+def write_content(data: bytes, name: str = "", extension: str = ""):
+    print("name", name, "extension", extension)
     path = str(Path(ARGS.root) / "txt")
     if not os.path.isdir(path):
         print(f"Missing '{path}' folder.")
         return None
     with tempfile.NamedTemporaryFile(
-            mode="wb",
-            dir=path,
-            suffix=".html",
-            delete=False
+        mode="wb",
+        dir=path,
+        suffix=f".{extension}" if extension else None,
+        delete=False,
     ) as output_file:
         output_file.write(data)
         os.chmod(output_file.name, 0o644)
-        return output_file.name
+        file_name = output_file.name
+    if name:
+        new_file_name = os.path.join(os.path.dirname(file_name), name)
+        try:
+            os.rename(file_name, name)
+            file_name = new_file_name
+        except OSError as exc:
+            print(f"Failed to give required name to the file: {exc}")
+            return None
+    return file_name
 
 
 def main():