xenu_nntp/lib/nntp/tiny/session.py

import re
import enum
import socket
import datetime
import fnmatch

from typing import Optional

from nntp.tiny.buffer    import LineBuffer, BufferOverflow
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()

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

    RE_NUM         = re.compile(r'^(\d+)$')
    RE_RANGE       = re.compile(r'^(\d+)-(\d+)$')
    RE_RANGE_LOWER = re.compile(r'^(\d+$)-$')
    RE_RANGE_UPPER = re.compile(r'^-(\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' % (NNTP_VERSION),
        'READER',
        'HDR',
        'NEWNEWS',
        'LIST ACTIVE NEWSGROUP OVERVIEW.FMT SUBSCRIPTIONS',
        'OVER MSGID'
    ]

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

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

        self.respond(ResponseCode.NNTP_GROUP_LISTING, summary)

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

        return self.end()

    def _cmd_list(self, *args):
        pass

    RE_DATE_SHORT = re.compile(r'^(\d{2})(\d{2})(\d{2})$')
    RE_DATE_LONG  = re.compile(r'^(\d{4})(\d{2})(\d{2})$')
    RE_TIME       = re.compile(r'^(\d{2})(\d{2})(\d{2})$')

    def _parse_date_time(self, datestr: str, timestr: str):
        yyyy, mm, dd = None, None, None,
        hh, MM, ss   = None, None, None
        
        match = self.RE_DATE_SHORT.match(datestr)
        if match:
            yy, mm, dd = map(int, match[1:3])

            if yy >= 70:
                yyyy = 1900 + yy
            else:
                yyyy = 2000 + yy

        match = self.RE_DATE_LONG.match(datestr)
        if match:
            yyyy, mm, dd = map(int, match[1:3])

        if yyyy is None:
            return

        match = self.RE_TIME.match(timestr)
        if match is None:
            return

        hh, mm, ss = map(int, match[1:3])

        return datetime.datetime(yyyy, mm, dd, hh, MM, ss)

    def _cmd_newnews(self, wildmat, datestr, timestr, *args):
        gmt = False

        if len(args) == 1:
            if args[0] == "GMT":
                gmt = True
            else:
                return self.send_response(ResponseCode.NNTP_SYNTAX_ERROR, "Only optional 'GMT' allowed")
        elif len(args) > 1:
            return self.send_response(ResponseCode.NNTP_SYNTAX_ERROR, "Too many arguments")

        timestamp = self._parse_date_time(datestr, timestr)

        if timestamp is None:
            return self.send_response(ResponseCode.NNTP_SYNTAX_ERROR, "Invalid date or time")

        self.respond(ResponseCode.NNTP_ARTICLE_LISTING_ID_FOLLOWS)

        sql = """
            select message_id from newsgroup_message where newsgroup_id = ? and created_on >= ?
        """

        for name in self.server.newsgroups:
            if fnmatch.fnmatch(name, wildmat):
                newsgroup = self.server.newsgroups[name]

                cr = self.db.execute(sql, (newsgroup.id, timestamp.isoformat()))

                for row in cr.each():
                    self.print(row[0])

        return self.end()

    def _cmd_newgroups(self, wildmat, datestr, timestr, *args):
        gmt = False

        if len(args) == 1:
            if args[0] == "GMT":
                gmt = True
            else:
                return self.send_response(ResponseCode.NNTP_SYNTAX_ERROR, "Only optional 'GMT' allowed")
        elif len(args) > 1:
            return self.send_response(ResponseCode.NNTP_SYNTAX_ERROR, "Too many arguments")

        pass

    COMMANDS = {
        'CAPABILITIES': _cmd_capabilities,
        'GROUP':        _cmd_group,
        'LISTGROUP':    _cmd_listgroup,
        'LIST':         _cmd_list,
        'NEWNEWS':      _cmd_newnews,
        'NEWSGROUPS':   _cmd_newgroups,
    }

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

        if line == '':
            return

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

        fn = self.COMMANDS.get(command.upper())

        if fn is None:
            return self.respond(ResponseCode.NNTP_COMMAND_UNKNOWN)

        return fn(self, *args)