You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
167 lines
4.7 KiB
167 lines
4.7 KiB
#!/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 = """\
|
|
<!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):
|
|
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 = "<pre>{}</pre>".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 = ""):
|
|
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()
|