server.py

import argparse
import socket
import logging
import threading


LOG = logging.getLogger(__name__)
"""Logger instance."""

DEBUG = False
"""Debug mode."""

IPC_TIMEOUT = 60
END_TCP_CONNECTION = b'\xff\xf4\xff\xfd\x06'
"""Dunno, investigate later."""

CLIENTS = {}
"""Client mapping (client_socket, client_address) -> Thread."""


def parse_args():
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("-h", "--host", default="localhost")
    parser.add_argument("-p", "--port", type=int, default=8080)
    parser.add_argument("-d", "--debug", action="store_true")

    return parser.parse_args()


def setup_logger(mode):
    logging.basicConfig(level=logging.DEBUG)
    #sh = logging.StreamHandler()
    #sh.setLevel(logging.DEBUG)
    #LOG.addHandler(sh)


def sigint(signum, frame):
    LOG.info("Caught signal num %s", signum)
    cleanup()


def setup_server(host, port):
    """Create server socket and setup it's opts.

    :param str host: server hostname
    :param int port: server port

    :return socket.socket: initialized server socket
    """
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
    server_socket.bind((host, port))
    server_socket.listen(5)

    return server_socket


def client_worker(sock, addr):

    def close():
        try:
            sock.close()
            del CLIENTS[sock, addr]
        except KeyError:
            LOG.warning("Client was not in registry.")
        except Exception:
            LOG.error("Socket closing failed.")


    while True:
        try:
            request = sock.recv(1024)
            LOG.debug("Got request: t=%s, s=%s", type(request), request)
            if request == END_TCP_CONNECTION:
                LOG.warning("Client closed connection.")
                close()
                return

            if request:
                request = request.decode()

            print(request)

            if "enough, pls" in request:
                LOG.info("Client wants to close connection.")
                sock.sendall(b"Ok, then")
                close()
                return

            sock.sendall(b"SRV OK")
        except Exception as e:
            LOG.exception("Something went wrong.")
            close()
            return


def serve_forever(host, port):
    server_socket = setup_server(host, port)
    LOG.info("Listening on %s:%s...", host, port)

    while True:
        client_socket, client_address = server_socket.accept()
        LOG.info("Got connection: {0}".format(client_address))
        client_socket.settimeout(60)

        for client in CLIENTS:
            client[0].sendall(b"New connected!")
        t = threading.Thread(target=client_worker,
                             args=(client_socket, client_address))
        t.start()
        CLIENTS[client_socket, client_address] = t


def cleanup():
    for t in CLIENTS.values():
        t.join(IPC_TIMEOUT)


if __name__ == "__main__":
    args = parse_args()
    setup_logger(args.debug)

    try:
        serve_forever(args.host, args.port)
    except KeyboardInterrupt as exc:
        LOG.info("Shutting down server...")
        cleanup()