xenu_nntp/lib/nntp/tiny/server.py

import re
import enum
import threading
import socket
import selectors
import ssl

from nntp.tiny.config    import Config, 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: Config):
        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.get('tls', 'cert'),
                                        config.get('tls', 'key'))

        self._init_newsgroups()

    def connect_to_db(self):
        return Database.connect(self.config.get('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 = None, None

        try:
            sock, addr = listener.accept()
        except ssl.SSLError as e:
            return

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

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

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

    @staticmethod
    def _is_ipv6(value: str):
        return value.find(':') >= 0

    @staticmethod
    def _is_ipv4(value: str):
        parts = value.split('.')

        if len(parts) > 4 or len(parts) == 0:
            return False

        for part in parts:
            if not part.isdecimal():
                return False

            num = int(part)

            if num > 255:
                return False

        return True

    def run(self):
        hosts = re.split(r'\s*,\s*', self.config.get('listen', 'host'))
        port  = int(self.config.get('listen', 'port'))

        listeners = list()

        for host in hosts:
            if Server._is_ipv6(host):
                listeners.append(self.listen(host, port, socket.AF_INET6))
            elif Server._is_ipv4(host):
                listeners.append(self.listen(host, port, socket.AF_INET))
            else:
                for af in (socket.AF_INET, socket.AF_INET6):
                    listeners.append(self.listen(host, port, af))

        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)