xenu_nntp/lib/nntp/tiny/server.py

import enum
import threading
import socket
import selectors
import ssl

from configparser import ConfigParser

from nntp.tiny.config    import ConfigException
from nntp.tiny.db        import Database
from nntp.tiny.newsgroup import Newsgroup
from nntp.tiny.session   import Session

class ServerCapability(enum.Flag):
    NONE = 0
    AUTH = enum.auto()
    POST = enum.auto()

class Server():
    def __init__(self, config: ConfigParser):
        self.config       = config
        self.capabilities = ServerCapability.NONE
        self.newsgroups   = dict()
        self.sslctx       = None

        if config['listen'].get('tls', 'no') == 'yes':
            self.sslctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
            self.sslctx.load_cert_chain(config['tls']['cert'],
                                        config['tls']['key'])

        self._init_newsgroups()

    def connect_to_db(self):
        return Database.connect(self.config['database']['path'])

    def _init_newsgroups(self):
        db = self.connect_to_db()

        for newsgroup in db.query(Newsgroup).each():
            self.newsgroups[newsgroup.name.casefold()] = newsgroup

    def listen(self, host: str, port: int, af: int):
        listener = socket.socket(af, socket.SOCK_STREAM)
        listener.bind((host, port))
        listener.listen()

        if self.sslctx:
            return self.sslctx.wrap_socket(listener, server_side=True)

        return listener

    def accept(self, listener):
        sock, addr = listener.accept()

        def spawn():
            session = Session(self, sock)

            try:
                session.handle()
            except (ssl.SSLError, ssl.SSLEOFError) as e:
                pass

        thread = threading.Thread(target=spawn)
        thread.start()

    def run(self):
        port = int(self.config['listen']['port'])

        listeners = list()

        if self.config.has_option('listen', 'host_inet'):
            host = self.config.get('listen', 'host_inet')
            listeners.append(self.listen(host, port, socket.AF_INET))

        if self.config.has_option('listen', 'host_inet6'):
            host = self.config.get('listen', 'host_inet6')
            listeners.append(self.listen(host, port, socket.AF_INET6))

        if len(listeners) == 0:
            raise ConfigException("No listener hosts specified")

        sel = selectors.DefaultSelector()

        for listener in listeners:
            sel.register(listener, selectors.EVENT_READ)

        while True:
            events = sel.select()

            for key, ev in events:
                self.accept(key.fileobj)