diff --git a/lib/nntp/tiny/session.py b/lib/nntp/tiny/session.py index 1153878..4efa102 100644 --- a/lib/nntp/tiny/session.py +++ b/lib/nntp/tiny/session.py @@ -19,6 +19,11 @@ class SessionState(enum.Flag): 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', @@ -103,7 +108,8 @@ class Session(): self.sock: socket.socket = sock self.buf: LineBuffer = LineBuffer() - self.newsgroup: Optional[Newsgroup] = None + self.newsgroup: Optional[Newsgroup] = None + self.article_id: Optional[int] = None def readline(self): return self.buf.readline(self.sock) @@ -162,7 +168,8 @@ class Session(): self.respond(ResponseCode.NNTP_GROUP_LISTING, text) - self.newsgroup = newsgroup + self.newsgroup = newsgroup + self.article_id = None return @@ -400,6 +407,74 @@ class Session(): return self.end() + def _message_by_id(self, identifier: str): + cr = None + + if identifier[0] == '<': + cr = self.db.query(Message, { + 'newsgroup_id': self.newsgroup.id, + 'message_id': identifier + }) + else: + cr = self.db.query(Message, { + 'newsgroup_id': self.newsgroup.id, + 'id': int(identifier) + }) + + return cr.fetchone() + + 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]): + 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(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]): + return self._serve_message(MessagePart.HEAD, identifier) + + def _cmd_body(self, identifier: Optional[str]): + return self._serve_message(MessagePart.BODY, identifier) + + def _cmd_article(self, identifier: Optional[str]): + return self._serve_message(MessagePart.WHOLE, identifier) + COMMANDS = { 'CAPABILITIES': _cmd_capabilities, 'GROUP': _cmd_group, @@ -407,6 +482,9 @@ class Session(): 'LIST': _cmd_list, 'NEWNEWS': _cmd_newnews, 'NEWGROUPS': _cmd_newgroups, + 'HEAD': _cmd_head, + 'BODY': _cmd_body, + 'ARTICLE': _cmd_article } def handle(self):