feat: add imap proxy server (#225)

smtp_proxy_server/config.py

@@ -0,0 +1,21 @@
+import logging
+from pydantic_settings import BaseSettings
+
+logging.basicConfig(
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+)
+
+_logger = logging.getLogger(__name__)
+_logger.setLevel(logging.INFO)
+
+
+class Settings(BaseSettings):
+    proxy_url: str = "http://localhost:8787"
+    port: int = 8025
+    imap_port: int = 11143
+
+    class Config:
+        env_file = ".env"
+
+
+settings = Settings()

smtp_proxy_server/docker-compose.yaml

@@ -7,6 +7,8 @@ services:
     container_name: "smtp_proxy_server"
     ports:
       - "8025:8025"
+      - "11143:11143"
     environment:
       - proxy_url=https://temp-email-api.xxx.xxx
       - port=8025
+      - imap_port=11143

smtp_proxy_server/dockerfile

@@ -4,4 +4,4 @@ WORKDIR /app
 COPY requirements.txt /requirements.txt
 RUN python3 -m pip install -r /requirements.txt
 COPY . /app
-ENTRYPOINT [ "python3", "server.py" ]
+ENTRYPOINT [ "python3", "main.py" ]

smtp_proxy_server/imap_server.py

@@ -0,0 +1,199 @@
+import json
+import logging
+import requests
+
+from io import BytesIO
+from twisted.mail import imap4
+from zope.interface import implementer
+from twisted.cred.portal import Portal, IRealm
+from twisted.internet import protocol, reactor, defer
+from twisted.cred.checkers import ICredentialsChecker, IUsernamePassword
+
+from config import settings
+from parse_email import parse_email
+from models import EmailModel
+
+_logger = logging.getLogger(__name__)
+_logger.setLevel(logging.INFO)
+
+
+@implementer(imap4.IMessage)
+class SimpleMessage:
+
+    def __init__(self, uid=None, email_model: EmailModel = None):
+        self.uid = uid
+        self.email = email_model
+        self.subparts = self.email.subparts
+
+    def getUID(self):
+        return self.uid
+
+    def getHeaders(self, negate, *names):
+        self.got_headers = negate, names
+        return {
+            k.lower(): v
+            for k, v in self.email.headers.items()
+        }
+
+    def isMultipart(self):
+        return len(self.subparts) > 0
+
+    def getSubPart(self, part):
+        self.got_subpart = part
+        return SimpleMessage(email_model=self.subparts[part])
+
+    def getBodyFile(self):
+        return BytesIO(self.email.body.encode("utf-8"))
+
+    def getSize(self):
+        return self.email.size
+
+    def getFlags(self):
+        return ["\\Seen"]
+
+
+@implementer(imap4.IMailboxInfo, imap4.IMailbox)
+class SimpleMailbox:
+
+    def __init__(self, password):
+        self.password = password
+        self.listeners = []
+        self.addListener = self.listeners.append
+        self.removeListener = self.listeners.remove
+        self.message_count = 0
+
+    def getFlags(self):
+        return ["\\Seen"]
+
+    def getUIDValidity(self):
+        return 0
+
+    def getMessageCount(self):
+        return 2 ** 31 - 1
+
+    def getRecentCount(self):
+        return 0
+
+    def getUnseenCount(self):
+        return 0
+
+    def isWriteable(self):
+        return 0
+
+    def destroy(self):
+        pass
+
+    def getHierarchicalDelimiter(self):
+        return "/"
+
+    def requestStatus(self, names):
+        r = {}
+        if "MESSAGES" in names:
+            r["MESSAGES"] = self.getMessageCount()
+        if "RECENT" in names:
+            r["RECENT"] = self.getRecentCount()
+        if "UIDNEXT" in names:
+            r["UIDNEXT"] = self.getMessageCount() + 1
+        if "UIDVALIDITY" in names:
+            r["UIDVALIDITY"] = self.getUID()
+        if "UNSEEN" in names:
+            r["UNSEEN"] = self.getUnseenCount()
+        return defer.succeed(r)
+
+    def fetch(self, messages, uid):
+        start, end = messages.ranges[0]
+        start = max(start, 1)
+        if start > self.message_count:
+            return []
+        res = requests.get(
+            f"{settings.proxy_url}/api/mails?limit=20&offset={start - 1}", headers={
+                "Authorization": f"Bearer {self.password}",
+                "Content-Type": "application/json"
+            }
+        )
+        if res.status_code != 200:
+            _logger.error(
+                "Failed: "
+                f"code=[{res.status_code}] text=[{res.text}]"
+            )
+            raise Exception("Failed to fetch emails")
+        if res.json()["count"] > 0:
+            self.message_count = res.json()["count"]
+        return [
+            (start + uid, SimpleMessage(start + uid, parse_email(item["raw"])))
+            for uid, item in enumerate(reversed(res.json()["results"]))
+        ]
+
+    def getUID(self, message):
+        return message.uid
+
+    def store(self, messages, flags, mode, uid):
+        # IMailboxIMAP.store
+        pass
+
+
+class Account(imap4.MemoryAccount):
+
+    def __init__(self, user, password):
+        self.password = password
+        super().__init__(user)
+
+    def _emptyMailbox(self, name, id):
+        _logger.info(f"New mailbox: {name}, {id}")
+        if name != "INBOX":
+            raise imap4.NoSuchMailbox(name)
+        return SimpleMailbox(self.password)
+
+    def select(self, name, rw=1):
+        return imap4.MemoryAccount.select(self, name)
+
+
+class SimpleIMAPServer(imap4.IMAP4Server):
+    def __init__(self, factory):
+        imap4.IMAP4Server.__init__(self)
+        self.factory = factory
+
+    def lineReceived(self, line):
+        super().lineReceived(line)
+
+
+@implementer(IRealm)
+class SimpleRealm:
+    def requestAvatar(self, avatarId, mind, *interfaces):
+        res = json.loads(avatarId)
+        account = Account(res["username"], res["password"])
+        account.addMailbox("INBOX")
+        return imap4.IAccount, account, lambda: None
+
+
+class IMAPFactory(protocol.Factory):
+    def __init__(self, portal):
+        self.portal = portal
+
+    def buildProtocol(self, addr):
+        p = SimpleIMAPServer(self)
+        p.portal = self.portal
+        return p
+
+
+@implementer(ICredentialsChecker)
+class CustomChecker:
+    credentialInterfaces = (IUsernamePassword,)
+
+    def requestAvatarId(self, credentials):
+        return defer.succeed(json.dumps({
+            "username": credentials.username.decode(),
+            "password": credentials.password.decode(),
+        }))
+
+
+def start_imap_server():
+    _logger.info(f"Starting IMAP server on port {settings.imap_port}")
+    portal = Portal(SimpleRealm(), [CustomChecker()])
+    reactor.listenTCP(settings.imap_port, IMAPFactory(portal))
+    reactor.run()
+
+
+if __name__ == "__main__":
+    _logger.info(f"Starting server settings[{settings}]")
+    start_imap_server()

smtp_proxy_server/main.py

@@ -0,0 +1,24 @@
+import logging
+import multiprocessing
+
+from smtp_server import start_smtp_server
+from imap_server import start_imap_server
+from config import settings
+
+_logger = logging.getLogger(__name__)
+_logger.setLevel(logging.INFO)
+
+if __name__ == '__main__':
+    _logger.info(f"Starting server settings[{settings}]")
+    process_list = [
+        multiprocessing.Process(target=start_smtp_server, args=()),
+        multiprocessing.Process(target=start_imap_server, args=()),
+    ]
+    try:
+        for p in process_list:
+            p.start()
+        for p in process_list:
+            p.join()
+    except KeyboardInterrupt:
+        for p in process_list:
+            p.terminate()

smtp_proxy_server/models.py

@@ -0,0 +1,10 @@
+from typing import Dict, List
+from pydantic import BaseModel
+
+
+class EmailModel(BaseModel):
+    headers: Dict[str, str]
+    body: str
+    content_type: str
+    subparts: List["EmailModel"]
+    size: int

smtp_proxy_server/parse_email.py

@@ -0,0 +1,38 @@
+import email
+from email.message import Message
+import logging
+
+from models import EmailModel
+
+_logger = logging.getLogger(__name__)
+_logger.setLevel(logging.INFO)
+
+
+def get_email_model(msg: Message):
+    subparts = [
+        get_email_model(subpart)
+        for subpart in msg.get_payload()
+    ] if msg.is_multipart() else []
+    body = "" if msg.is_multipart() else msg._payload
+    return EmailModel(
+        headers={k: v for k, v in msg.items()},
+        body=body,
+        content_type=msg.get_content_type(),
+        size=len(body) + sum(subpart.size for subpart in subparts),
+        subparts=subparts,
+    )
+
+
+def parse_email(raw: str) -> EmailModel:
+    try:
+        msg = email.message_from_string(raw)
+        return get_email_model(msg)
+    except Exception as e:
+        _logger.error(f"Could not parse email: {e}")
+        return EmailModel(
+            headers={},
+            body="could not parse email",
+            content_type="text/plain",
+            size=len("could not parse email"),
+            subparts=[],
+        )

smtp_proxy_server/requirements.txt

@@ -1,3 +1,4 @@
 aiosmtpd==1.4.5
 pydantic-settings==2.2.1
 requests==2.31.0
+twisted==24.3.0

smtp_proxy_server/smtp_server.py

@@ -3,26 +3,15 @@ import logging
 import email
 import requests
 
-from pydantic_settings import BaseSettings
 from aiosmtpd.controller import Controller
 from aiosmtpd.smtp import SMTP, Session, Envelope, AuthResult, LoginPassword
 
-logging.basicConfig(
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-)
+from config import settings
 
 _logger = logging.getLogger(__name__)
 _logger.setLevel(logging.INFO)
 
 
-class Settings(BaseSettings):
-    proxy_url: str = "http://localhost:8787"
-    port: int = 8025
-
-    class Config:
-        env_file = ".env"
-
-
 class CustomSMTPHandler:
 
     def authenticator(self, server, session, envelope, mechanism, auth_data):
@@ -119,7 +108,6 @@ class CustomSMTPHandler:
         return '250 OK'
 
 
-settings = Settings()
 handler = CustomSMTPHandler()
 server = Controller(
     handler,
@@ -133,11 +121,11 @@ server = Controller(
 
 
 async def start():
-    _logger.info(f"Starting server settings[{settings}]")
+    _logger.info(f"Starting server on port {settings.port}")
     server.start()
 
 
-if __name__ == "__main__":
+def start_smtp_server():
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
     task = loop.create_task(start())
@@ -146,3 +134,8 @@ if __name__ == "__main__":
     except KeyboardInterrupt:
         _logger.info("Got KeyboardInterrupt, stopping")
         server.stop()
+
+
+if __name__ == "__main__":
+    _logger.info(f"Starting server settings[{settings}]")
+    start_smtp_server()