cn-py-serv/server/response.py

#
# CatNet Python Server
#
# Copyright (C) 2023-2024 John Solntsev
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
#

import logging


log = logging.getLogger(__name__)

STATUS_BY_CODE = {
    200: "OK",

    400: "Bad Request",
    404: "Not Found",
    405: "Method Not Allowed",
    413: "Payload Too Large",

    500: "Internal Server Error",
    505: "HTTP Version Not Supported",
}


class Response:
    def __iter__(self):
        return self

    def __init__(self, status_code=200, additional_headers={}, data=None):
        self.status_code = status_code
        self.additional_headers = additional_headers
        self.data = data
        self.first = True

    def __next__(self):
        if self.first:
            log.debug("First packet generation of response")

            resp = "HTTP/1.1 {} {}\r\n"
            resp += "Server: cnserv\r\n"
            resp += "Connection: keep-alive\r\n"

            for key, val in self.additional_headers.items():
                resp += "{}: {}\r\n".format(key, val)

            if type(self.data) == bytes:
                resp += "Content-Length: {}\r\n".format(len(self.data))
            elif hasattr(self.data, "__iter__"):
                resp += "Content-Length: {}\r\n".format(next(self.data))

            resp += "\r\n"
            resp = resp.format(self.status_code, STATUS_BY_CODE[self.status_code])
            resp = resp.encode("UTF-8")

            if type(self.data) == bytes:
                resp += self.data

            self.first = False

            return resp

        else:
            if type(self.data) == bytes:
                raise StopIteration

            elif hasattr(self.data, "__iter__"):
                log.debug("Data is iterator, just generate next part")
                try:
                    return next(self.data)
                except StopIteration:
                    raise StopIteration