xenu_nntp/lib/nntp/tiny/session.py

import enum
import socket
import re

from typing import Optional

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 __dict__(self):
        if self.id is not None:
            return {'id': self.id}

        if self.min is not None and self.max is None:
            return { 'min(id)': self.min }
        elif self.min is not None and self.max is not None:
            return { 'min(id)': self.min, 'max(id)': self.max }
        elif self.min is None and self.max is not None:
            return { 'max(id)': 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.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.server.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 _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)

        values = {
            'newsgroup_id': newsgroup.id
        }

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

            values.update(dict(msgrange))

        cr = self.server.db.query(Message, values)

        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())