#!/usr/bin/env python3
"""shrlok server: receive text/files from a socket, put them on a Web server.
The server expects messages with the following format to come through the
socket:
1. a length, as ASCII digits,
2. a null byte,
3. a JSON object containing at least the "type" key with a known value,
4. another null byte,
5. the text or file itself.
The length is of the JSON object + the null char + the content itself.
Example, with \\0 representing a null byte:
28\\0{"type":"txt"}\\0hello shrlok!
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.
"""
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):
data = self.receive_input()
# Extract header.
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
file_name = None
if header.get("type") == "txt":
file_name = write_text(data, title=header.get("title", ""))
elif header.get("type") == "raw":
file_name = write_content(
data,
name=header.get("name", ""),
extension=header.get("ext", ""),
)
else:
print("Unknown type.")
if file_name is None:
return
print(f"{len(data)} bytes — {header} — '{file_name}'.")
try:
self.request.sendall(file_name.encode())
except BrokenPipeError:
print("Broken pipe.")
def receive_input(self):
length = None
data = b""
while True:
chunk = self.request.recv(4096)
if not chunk:
break
data += chunk
if length is None:
try:
first_zero = data.index(b"\0")
except ValueError:
return b""
length_data, data = data[:first_zero], data[first_zero + 1:]
length = int(length_data.decode())
if len(data) >= length: # retrieval completed
break
return data
def write_text(data: bytes, title: str = "", name: str = ""):
content = "{}".format(data.decode(errors="replace"))
html = HTML_TEMPLATE.format(title=title or "", content=content).encode()
return write_content(html, extension="html")
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=f".{extension}" if extension else None,
delete=False,
) as output_file:
output_file.write(data)
os.chmod(output_file.name, 0o644)
file_name = output_file.name
if name:
old_file_name = file_name
file_name = os.path.join(os.path.dirname(old_file_name), name)
try:
os.rename(old_file_name, file_name)
except OSError as exc:
print(f"Failed to give required name to the file: {exc}")
return None
return 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:
os.chmod(socket_path, 0o664)
server.serve_forever()
except KeyboardInterrupt:
print("Stopping server.")
finally:
os.unlink(socket_path)
if __name__ == "__main__":
main()