xenu_nntp/lib/nntp/tiny/server.py

import enum
import socket
import ssl
import threading

from configparser import ConfigParser

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 run(self):
        host = self.config['listen']['host']
        port = int(self.config['listen']['port'])

        listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        listener.bind((host, port))
        listener.listen()

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

        while True:
            try:
                sock, addr = listener.accept()

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

                thread = threading.Thread(target=spawn)
                thread.start()
            except ssl.SSLEOFError as e:
                pass

        listener.close()