This repository has been archived on 2024-08-20. You can view files and clone it, but cannot push or open issues or pull requests.
Shrlok/shrlok/shrlok.py
2022-07-04 21:52:48 +02:00

92 lines
2.2 KiB
Python

#!/usr/bin/env python3
import argparse
import json
import os
import socketserver
import tempfile
from pathlib import Path
ARGS = None
HTML_TEMPLATE = """\
<!doctype html>
<html>
<head>
<meta charset="UTF-8">
<title>{title}</title>
<style>
body {{ max-width: 40em; }}
</style>
</head>
<body>
{content}
</body>
</html>
"""
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
print(f"{len(data)} bytes — {header}'{file_name}'.")
self.request.sendall(file_name.encode())
def write_text(data: bytes, title=None):
content = "<pre>{}</pre>".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):
os.makedirs(path)
with tempfile.NamedTemporaryFile(
mode="wt",
dir=path,
suffix=".html",
delete=False
) as output_file:
output_file.write(html)
return output_file.name
def main():
argp = argparse.ArgumentParser()
argp.add_argument("root", help="root path where to put files")
global ARGS
ARGS = argp.parse_args()
uid = os.getuid()
socket_path = f"/run/user/{uid}/shr.sock"
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()