xenu_nntp/lib/nntp/tiny/session.py

import re
import enum
import socket
import datetime
import fnmatch
import traceback

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 MessagePart(enum.Enum):
    HEAD  = 1
    BODY  = enum.auto()
    WHOLE = 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 = int(match[1])
            return obj

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

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

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

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

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

    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
        self.article_id: Optional[int]       = 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, 'utf-8'))

    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
        self.article_id = None

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

        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 _newsgroup_summary(self, newsgroup: Newsgroup):
        sql = """
            select
                min(id),
                max(id)
            from
                newsgroup_message
            where
                newsgroup_id = ?
        """

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

        return {
            'low':   row[0],
            'high':  row[1],
            'perms': 'n'
        }


    def print_newsgroup(self, newsgroup: Newsgroup):
        summary = self._newsgroup_summary(newsgroup)

        return self.print("%s %d %d %s" % (
            newsgroup.name,
            summary['low'],
            summary['high'],
            summary['perms']
        ))

    def _cmd_list_newsgroups(self):
        self.respond(ResponseCode.NNTP_INFORMATION_FOLLOWS)

        for name in self.server.newsgroups:
            newsgroup = self.server.newsgroups[name]

            self.print_newsgroup(newsgroup)

        return self.end()

    def _newsgroup_last_active(self, newsgroup: Newsgroup):
        sql = """
            select
                max(created_on)
            from
                newsgroup_message
            where
                newsgroup_id = ?
        """

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

        if row is None:
            return

        return datetime.datetime.fromisoformat(row[0])

    def _cmd_list_active(self):
        now = datetime.datetime.now(datetime.UTC)

        self.respond(ResponseCode.NNTP_INFORMATION_FOLLOWS)

        for name in self.server.newsgroups:
            newsgroup   = self.server.newsgroups[name]
            last_active = self._newsgroup_last_active(newsgroup)

            if now - last_active < datetime.timedelta(days=1):
                self.print_newsgroup(newsgroup)

        return self.end()

    LIST_SUBCOMMANDS = {
        'NEWSGROUPS': _cmd_list_newsgroups,
        'ACTIVE':     _cmd_list_active,
    }

    def _cmd_list(self, *args):
        if len(args) == 0:
            return self.respond(ResponseCode.NNTP_SYNTAX_ERROR, "No subcommand provided")

        subcmd, *subargs = args

        fn = self.LIST_SUBCOMMANDS.get(subcmd.upper())

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

        return fn(self, *subargs)

    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.respond(ResponseCode.NNTP_SYNTAX_ERROR, "Only optional 'GMT' allowed")
        elif len(args) > 1:
            return self.respond(ResponseCode.NNTP_SYNTAX_ERROR, "Too many arguments")

        self.respond(ResponseCode.NNTP_GROUPS_NEW_FOLLOW)

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

                self.print_newsgroup(newsgroup)

        return self.end()

    def _message_by_id(self, identifier: str):
        if identifier[0] == '<':
            return self.db.get(Message, {
                'newsgroup_id': self.newsgroup.id,
                'message_id':   identifier
            })
        else:
            return self.db.get(Message, {
                'newsgroup_id': self.newsgroup.id,
                'id':           int(identifier)
            })

    def _send_message_headers(self, message: Message):
        for name in message.headers:
            self.print("%s: %s" % (
                name, message.headers[name]
            ))

    def _serve_message(self, part: MessagePart, identifier: Optional[str]=None):
        if self.newsgroup is None:
            return self.respond(ResponseCode.NNTP_NEWSGROUP_NOT_SELECTED)

        message = None

        if self.article_id is None:
            if identifier is None:
                return self.respond(ResponseCode.NNTP_ARTICLE_INVALID_NUMBER)
            else:
                message = self._message_by_id(identifier)
                self.article_id = message.id
        else:
            if identifier is None:
                message = self._message_by_id(str(self.article_id))
            else:
                message = self._message_by_id(identifier)
                self.article_id = message.id

        text = "%d %s" % (
            message.id,
            message.message_id
        )

        self.respond(ResponseCode.NNTP_ARTICLE_LISTING, text)

        if part is MessagePart.HEAD or part is MessagePart.WHOLE:
            self._send_message_headers(message)

        if part is MessagePart.WHOLE:
            self.print('')

        if part is MessagePart.BODY or part is MessagePart.WHOLE:
            self.print(message.body)

        return self.end()

    def _cmd_head(self, identifier: Optional[str]=None):
        return self._serve_message(MessagePart.HEAD, identifier)

    def _cmd_body(self, identifier: Optional[str]=None):
        return self._serve_message(MessagePart.BODY, identifier)

    def _cmd_article(self, identifier: Optional[str]=None):
        return self._serve_message(MessagePart.WHOLE, identifier)

    def _send_message_header(self, message: Message, name: str):
        return self.print("%d %s" % (
            message.id, message.headers.get(name, '')
        ))

    def _cmd_hdr(self, name: str, msg: Optional[str]=None):
        if self.newsgroup is None:
            return self.respond(ResponseCode.NNTP_NEWSGROUP_NOT_SELECTED)

        if msg is None:
            if self.article_id is None:
                return self.respond(ResponseCode.NNTP_ARTICLE_INVALID_NUMBER)

            message = self.db.get(Message, self.article_id)

            self.respond(ResponseCode.NNTP_HEADERS_FOLLOW)
            self._send_message_header(message, name)
        else:
            msgrange = MessageRange.parse(msg)

            sql  = f"select * from {Message.name} where "
            sql += " newsgroup_id = ? and " + msgrange.where()

            cr = self.db.query_sql(Message, sql, (self.newsgroup.id,))

            self.respond(ResponseCode.NNTP_HEADERS_FOLLOW)

            for message in cr.each():
                self._send_message_header(message, name)

        return self.end()

    def _cmd_stat(self, identifier: Optional[str]=None):
        if self.newsgroup is None:
            return self.respond(ResponseCode.NNTP_NEWSGROUP_NOT_SELECTED)

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

        row = None

        if identifier is None:
            if self.article_id is None:
                return self.respond(ResponseCode.NNTP_ARTICLE_INVALID_NUMBER)

            identifier = str(self.article_id)

        if identifier[0] == '<':
            cr = self.db.execute(sql + " and message_id = ?",
                                 (self.newsgroup.id, identifier))

            row = cr.fetchone()

            if row is None:
                return self.respond(ResponseCode.NNTP_ARTICLE_NOT_FOUND_ID)
        else:
            cr = self.db.execute(sql + " and id = ?",
                                 (self.newsgroup.id, int(identifier)))

            row = cr.fetchone()

            if row is None:
                return self.respond(ResponseCode.NNTP_ARTICLE_NOT_FOUND_NUM)

        text = "%d %s" % (row[0], row[1])

        if self.article_id is None:
            self.article_id = int(row[0])

        return self.respond(ResponseCode.NNTP_ARTICLE_STAT_RESPONSE, text)

    def _cmd_date(self):
        timestamp = datetime.datetime.now(datetime.UTC)

        return self.respond(ResponseCode.NNTP_DATE,
                            timestamp.strftime("%Y%m%d%H%M%S"))

    COMMANDS = {
        'CAPABILITIES': _cmd_capabilities,
        'GROUP':        _cmd_group,
        'LISTGROUP':    _cmd_listgroup,
        'LIST':         _cmd_list,
        'NEWNEWS':      _cmd_newnews,
        'NEWGROUPS':    _cmd_newgroups,
        'HEAD':         _cmd_head,
        'BODY':         _cmd_body,
        'ARTICLE':      _cmd_article,
        'HDR':          _cmd_hdr,
        'STAT':         _cmd_stat,
        'DATE':         _cmd_date,
    }

    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)

        try:
            return fn(self, *args)
        except TypeError as e:
            traceback.print_exception(e)
            return self.respond(ResponseCode.NNTP_SYNTAX_ERROR)
        except Exception as e:
            traceback.print_exception(e)
            return self.respond(ResponseCode.NNTP_COMMAND_UNAVAILABLE)