Implement HEAD, BODY, ARTICLE

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',
 
@@ -104,6 +109,7 @@ class Session():
         self.buf: LineBuffer     = LineBuffer()
 
         self.newsgroup:  Optional[Newsgroup] = None
+        self.article_id: Optional[int]       = None
 
     def readline(self):
         return self.buf.readline(self.sock)
@@ -163,6 +169,7 @@ class Session():
         self.respond(ResponseCode.NNTP_GROUP_LISTING, text)
 
         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):