xenu_nntp/lib/nntp/tiny/session.py

import enum
import socket
import re

from typing import Optional

from nntp.tiny.db        import Database
from nntp.tiny.server    import Server, ServerCapability
from nntp.tiny.response  import Response, ResponseCode
from nntp.tiny.newsgroup import Newsgroup
from nntp.tiny.message   import Message

class SessionState(enum.Flag):
    NONE      = 0
    AUTH_OK   = enum.auto()
    AUTH_POST = enum.auto()

from nntp.tiny.buffer import LineBuffer, BufferOverflow

class MessageRange():
    __slots__ = 'id', 'min', 'max',

    RE_NUM         = re.compile('^(\d+)$')
    RE_RANGE       = re.compile('^(\d+)-(\d+)$')
    RE_RANGE_LOWER = re.compile('^(\d+$)-$')
    RE_RANGE_UPPER = re.compile('^-(\d+$)$')

    def __init__(self):
        self.id:  int = None
        self.min: int = None
        self.max: int = None

    def __str__(self):
        if self.id is not None:
            return str(self.id)

        if self.min is not None and self.max is None:
            return "%d-" % (self.min)
        elif self.min is not None and self.max is not None:
            return "%d-%d" % (self.min, self.max)
        elif self.min is None and self.max is not None:
            return "-%d" % (self.max)

        return "?"

    def where(self):
        if self.id is not None:
            return "id = %d" % (self.id)

        if self.min is not None and self.max is None:
            return "id >= %d" % (self.min)
        elif self.min is not None and self.max is not None:
            return "id >= %d and id <= %d" % (self.min, self.max)
        elif self.min is None and self.max is not None:
            return "id <= %d" % (self.max)

    @staticmethod
    def parse(r: str):
        match = __class__.RE_NUM.match(r)
        if match:
            obj = __class__()
            obj.id = match[1]
            return obj

        match = __class__.RE_RANGE.match(r)
        if match:
            obj = __class__()
            obj.min = match[1]
            obj.max = match[2]
            return obj

        match = __class__.RE_RANGE_LOWER.match(r)
        if match:
            obj = __class__()
            obj.min = match[1]
            return obj

        match = __class__.RE_RANGE_UPPER.match(r)
        if match:
            obj = __class__()
            obj.max = match[1]
            return obj

class Session():
    RE_SPLIT = re.compile(r'\s+')

    NNTP_VERSION = 2

    NNTP_CAPABILITIES = [
        'VERSION %d' % (self.NNTP_VERSION),
        'READER',
        'HDR',
        'NEWNEWS',
        'LIST ACTIVE NEWSGROUP OVERVIEW.FMT SUBSCRIPTIONS',
        'OVER MSGID'
    ]

    COMMANDS = {
        'capabilities': 
    }
    
    def __init__(self, server: Server, sock: socket.socket):
        self.server: Server      = server
        self.db: Database        = server.db
        self.state: SessionState = SessionState.NONE
        self.sock: socket.socket = sock
        self.buf: LineBuffer     = LineBuffer()

        self.newsgroup: Optional[Newsgroup] = None

    def readline(self):
        return self.buf.readline(self.sock)

    def print(self, text: str, end: str="\r\n"):
        return self.sock.send(bytes(text + end, 'ascii'))

    def end(self):
        return self.print('.')

    def respond(self, code: ResponseCode, message: str=None, body=None):
        response = Response(code, message, body)

        return self.print(str(response))

    def _cmd_capabilities(self, *args):
        self.respond(ResponseCode.NNTP_CAPABILITIES_FOLLOW)

        if self.state & SessionState.AUTH_POST:
            self.print('POST')

        if self.state & SessionState.AUTH_OK:
            self.print('AUTHUSER INFO')

        for item in self.NNTP_CAPABILITIES:
            self.print(item)

        self.end()

    def _cmd_group(self, name: str):
        if name not in self.server.newsgroups:
            return self.respond(ResponseCode.NNTP_NEWSGROUP_NOT_FOUND)

        newsgroup = self.server.newsgroups[name]

        sql = """
            select
                count(id),
                min(id),
                max(id)
            from
                newsgroup_message
            where
                newsgroup_id = ?
        """

        cr  = self.db.execute(sql, (newsgroup.id))
        row = cr.fetchone()

        text = "%d %d %d %s" % (
            row[0],
            row[1],
            row[2],
            newsgroup.name
        )

        self.respond(ResponseCode.NNTP_GROUP_LISTING, text)

        self.newsgroup = newsgroup

        return

    def _newsgroup_summary(self, newsgroup: Newsgroup) -> str:
        sql = """
            select
                count(id),
                min(id),
                max(id)
            from
                newsgroup_message
            where
                newsgroup_id = ?
        """

        cr  = self.db.execute(sql, (newsgroup.id))
        row = cr.fetchone()

        return "%d %d %d %s" % (
            row[0],
            row[1],
            row[2],
            newsgroup.name
        )

    def _cmd_listgroup(self, *args):
        newsgroup = self.newsgroup
        
        if len(args) == 0 and newsgroup is None:
            return self.respond(ResponseCode.NNTP_NEWSGROUP_NOT_SELECTED)
        elif len(args) > 0:
            newsgroup = self.server.newsgroups.get(args[0])

            if newsgroup is None:
                return self.respond(ResponseCode.NNTP_NEWSGROUP_NOT_FOUND)

        sql = """
            select
                id
            from
                newsgroup_message
            where
                newsgroup_id = ?
        """

        if len(args) > 1:
            msgrange = MessageRange.parse(args[1])

            sql += " and " msgrange.clause()

        text = self._newsgroup_summary(newsgroup)
        cr   = self.db.execute(sql, (newsgroup.id))

        self.respond(ResponseCode.NNTP_GROUP_LISTING, text)

        for message in cr.each():
            self.print(str(message.id))

        return self.end()

    def handle(self):
        line = self.readline()

        if line == '':
            return

        command, args = self.RE_SPLIT.split(line.rstrip())