#!/usr/bin/env python3
import argparse
import json
import os
import socketserver
import tempfile
from pathlib import Path
ARGS = None
HTML_TEMPLATE = """\
{title}
{content}
"""
class Handler(socketserver.StreamRequestHandler):
def handle(self):
fragments = []
while (chunk := self.request.recv(4096)):
fragments.append(chunk)
data = b"".join(fragments)
try:
first_zero = data.index(b"\0")
header_data, data = data[:first_zero], data[first_zero + 1:]
header = json.loads(header_data.decode())
except (ValueError, IndexError):
print("Bad header.")
return
if header.get("type") == "txt":
file_name = write_text(data, title=header.get("title"))
else:
print("Unknown type.")
return
if not file_name:
return
print(f"{len(data)} bytes — {header} — '{file_name}'.")
self.request.sendall(file_name.encode())
def write_text(data: bytes, title=None):
content = "{}".format(data.decode(errors="replace"))
html = HTML_TEMPLATE.format(title=title or "?", content=content)
path = str(Path(ARGS.root) / "txt")
if not os.path.isdir(path):
print(f"Missing '{path}' folder.")
return None
with tempfile.NamedTemporaryFile(
mode="wt",
dir=path,
suffix=".html",
delete=False
) as output_file:
output_file.write(html)
os.chmod(output_file.name, 0o644)
return output_file.name
def main():
argp = argparse.ArgumentParser(description="Share stuff through a socket.")
argp.add_argument("root", help="root path where to put files")
argp.add_argument("-s", "--socket", help="socket path")
global ARGS
ARGS = argp.parse_args()
socket_path = ARGS.socket
if not socket_path:
runtime_dir = os.environ.get("RUNTIME_DIRECTORY")
if not runtime_dir:
exit("No socket path nor runtime directory specified.")
socket_path = os.path.join(runtime_dir, "shr.sock")
print("Socket path:", socket_path)
if os.path.exists(socket_path):
os.unlink(socket_path)
try:
with socketserver.UnixStreamServer(socket_path, Handler) as server:
server.serve_forever()
except KeyboardInterrupt:
print("Stopping server.")
finally:
os.unlink(socket_path)
if __name__ == "__main__":
main()