From f6100c4e2f3c8e39eaa041603af968f8eadc9c5b Mon Sep 17 00:00:00 2001
From: XANTRONIX Development <dev@xantronix.com>
Date: Tue, 26 Nov 2024 15:49:54 -0500
Subject: [PATCH] More refactoring to eliminate duplicate code

---
 lib/nntp/tiny/response.py |  4 +-
 lib/nntp/tiny/session.py  | 98 ++++++++++++++++++---------------------
 2 files changed, 46 insertions(+), 56 deletions(-)

diff --git a/lib/nntp/tiny/response.py b/lib/nntp/tiny/response.py
index 4e730fb..7ac9126 100644
--- a/lib/nntp/tiny/response.py
+++ b/lib/nntp/tiny/response.py
@@ -29,7 +29,7 @@ class ResponseCode(enum.Enum):
     NNTP_ARTICLE_INVALID_NUMBER        = 420
     NNTP_ARTICLE_NO_NEXT               = 421
     NNTP_ARTICLE_NO_PREVIOUS           = 422
-    NNTP_ARTICLE_NOT_FOUND_RANGE       = 423
+    NNTP_ARTICLE_NOT_FOUND_NUM         = 423
     NNTP_ARTICLE_NOT_FOUND_ID          = 430
     NNTP_POST_PROHIBITED               = 440
     NNTP_POST_FAILED                   = 441
@@ -61,7 +61,7 @@ class ResponseCode(enum.Enum):
             420: "Current article number is invalid",
             421: "No next article in this group",
             422: "No previous article in this group",
-            423: "No articles found within range",
+            423: "No article found by that message number",
             430: "No article found by that message ID",
             440: "Posting prohibited",
             441: "Posting failed",
diff --git a/lib/nntp/tiny/session.py b/lib/nntp/tiny/session.py
index 5c2ce24..bb92dda 100644
--- a/lib/nntp/tiny/session.py
+++ b/lib/nntp/tiny/session.py
@@ -528,7 +528,7 @@ class Session():
                 self.respond(ResponseCode.NNTP_ARTICLE_INVALID_NUMBER)
                 return
 
-            message = self.db.get(Message, self.article_id)
+            message = self.db.get(Message, {'id': str(self.article_id)})
 
             if message is None:
                 self.respond(ResponseCode.NNTP_ARTICLE_INVALID_NUMBER)
@@ -581,24 +581,47 @@ class Session():
                 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)
+    def _message_by_id(self, identifier: Optional[str]=None):
+        if identifier is None:
+            if self.newsgroup is None:
+                self.respond(ResponseCode.NNTP_NEWSGROUP_NOT_SELECTED)
+                return
 
-        message = None
+            if self.article_id is None:
+                self.respond(ResponseCode.NNTP_ARTICLE_INVALID_NUMBER)
+                return
 
-        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
+            message = self.db.get(Message, {'id': self.article_id})
+
+            if message is None:
+                self.respond(ResponseCode.NNTP_ARTICLE_NOT_FOUND_NUM)
+                return
+
+            return message
+        elif identifier[0] == '<':
+            message = self.db.query(Message, {
+                'message_id': identifier
+            }).fetchone()
+
+            if message is None:
+                self.respond(ResponseCode.NNTP_ARTICLE_NOT_FOUND_ID)
+                return
+
+            return message
         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
+            message = self.db.get(Message, {'id': int(identifier)})
+
+            if message is None:
+                self.respond(ResponseCode.NNTP_ARTICLE_NOT_FOUND_NUM)
+                return
+
+            return message
+
+    def _serve_message(self, part: MessagePart, identifier: Optional[str]=None):
+        message = self._message_by_id(identifier)
+
+        if message is None:
+            return
 
         text = "%d %s" % (
             message.id,
@@ -655,47 +678,14 @@ class Session():
             self.print('|'.join(overview))
 
     def _cmd_stat(self, identifier: Optional[str]=None):
-        if self.newsgroup is None:
-            return self.respond(ResponseCode.NNTP_NEWSGROUP_NOT_SELECTED)
+        message = self._message_by_id(identifier)
 
-        sql = """
-            select
-                id, message_id
-            from
-                newsgroup_message
-            where
-                newsgroup_id = ?
-        """
+        if message is None:
+            return
 
-        row = None
+        text = "%d %s" % (message.id, message.message_id)
 
-        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_RANGE)
-
-        text = "%d %s" % (row[0], row[1])
-
-        if self.article_id is None:
-            self.article_id = int(row[0])
+        self.article_id = message.id
 
         return self.respond(ResponseCode.NNTP_ARTICLE_STAT_RESPONSE, text)