sparty/sparty/sparty.py

#!/usr/bin/env python3
# vim: tabstop=4 shiftwidth=4 expandtab

import socket
import urllib.parse

class Response():
    def __init__(self, socket: socket.socket) -> "Response":
        self._socket = socket
        self._filehandle = self._socket.makefile(mode = "rb")

        try:
            self.responsecode, self.meta = self._filehandle.readline(1027).split(maxsplit=1)
            self.responsecode = int(self.responsecode)
            self.meta = self.meta.strip().decode("ASCII")
        except:
            self.discard()
            raise RuntimeError("Received malformed header from spartan server")

    def read(self, bufsize: int = 4096) -> "bytes":
        return self._filehandle.read(bufsize)

    def discard(self) -> None:
        self._filehandle.close()
        self._socket.close()

def request(url: str = "", message: bytes = "") -> "Response":
    url = f"spartan://{url}" if not url.startswith("spartan://") else url
    parsed = urllib.parse.urlparse(url)
    data = str.encode(parsed.query) or message
    sock = socket.create_connection((parsed.hostname, parsed.port or 300))
    sock.sendall(f"{parsed.hostname} {parsed.path or '/'} {len(data)}\r\n".encode("ASCII"))
    if data:
        sock.sendall(data)

    return Response(sock)