From e7c5f60fbe5262bacf67030636980926323bba9a Mon Sep 17 00:00:00 2001
From: Shihaam Abdul Rahman <shihaam@shihaam.me>
Date: Sat, 21 Sep 2024 05:45:26 +0500
Subject: [PATCH] add support for ntfy

---
 .build/Dockerfile |  2 +-
 .env.example      |  7 ++++--
 relay.py          | 60 +++++++++++++++++++++++++++++++++--------------
 3 files changed, 48 insertions(+), 21 deletions(-)

diff --git a/.build/Dockerfile b/.build/Dockerfile
index b677b7e..bf73ef9 100644
--- a/.build/Dockerfile
+++ b/.build/Dockerfile
@@ -4,6 +4,6 @@ WORKDIR /app
 
 COPY relay.py .
 
-RUN pip install --no-cache-dir aiosmtpd python-telegram-bot
+RUN pip install --no-cache-dir aiosmtpd python-telegram-bot httpx
 
 CMD ["python", "relay.py"]
diff --git a/.env.example b/.env.example
index 4b57497..6dbb759 100644
--- a/.env.example
+++ b/.env.example
@@ -1,2 +1,5 @@
-TELEGRAM_API_KEY=
-TELEGRAM_CHAT_ID=
+TELEGRAM_API_KEY=''
+TELEGRAM_CHAT_ID=''
+
+NTFY_TOKEN='your_ntfy_token_here'
+NTFY_URL='https://ntfy.sh'
diff --git a/relay.py b/relay.py
index 1fd62d1..2ebe104 100644
--- a/relay.py
+++ b/relay.py
@@ -5,10 +5,13 @@ from email.parser import BytesParser
 from email.policy import default
 import telegram
 import re
+import httpx
 
-# Telegram configuration
+# Configuration
 TELEGRAM_API_KEY = os.getenv('TELEGRAM_API_KEY')
 DEFAULT_CHAT_ID = os.getenv('TELEGRAM_CHAT_ID')  # Fallback chat ID
+NTFY_TOKEN = os.getenv('NTFY_TOKEN')
+NTFY_URL = os.getenv('NTFY_URL')
 
 # SMTP configuration
 SMTP_HOST = '0.0.0.0'  # Listen on all interfaces
@@ -25,29 +28,50 @@ class SMTPHandler:
         
         subject = message.get('subject', 'No subject')
         body = message.get_body(preferencelist=('plain', 'html')).get_content()
-
-        telegram_message = f"Subject: {subject}\n\n{body}"
+        full_message = f"Subject: {subject}\n\n{body}"
         
-        # Extract chat ID from the recipient email address
-        chat_id = self.extract_chat_id(envelope.rcpt_tos[0])
+        # Extract destination from the recipient email address
+        destination, service = self.extract_destination(envelope.rcpt_tos[0])
         
-        bot = telegram.Bot(TELEGRAM_API_KEY)
-        try:
-            await bot.send_message(chat_id=chat_id, text=telegram_message)
-            print(f"Message sent to Telegram chat ID: {chat_id}")
-        except telegram.error.BadRequest as e:
-            print(f"Failed to send message to chat ID: {chat_id}. Error: {str(e)}")
-            print("Attempting to send to default chat ID.")
-            await bot.send_message(chat_id=DEFAULT_CHAT_ID, text=telegram_message)
+        if service == 'telegram':
+            await self.send_telegram(destination, full_message)
+        elif service == 'ntfy':
+            await self.send_ntfy(destination, subject, body)
+        else:
+            print(f"Invalid service: {service}. Message not sent.")
 
         return '250 Message accepted for delivery'
 
-    def extract_chat_id(self, email):
-        match = re.match(r'(-?\d+)@telegram-relay\.local', email)
+    def extract_destination(self, email):
+        match = re.match(r'(.+)@(telegram|ntfy)$', email)
         if match:
-            return match.group(1)
-        print(f"Invalid email format: {email}. Using default chat ID.")
-        return DEFAULT_CHAT_ID
+            return match.group(1), match.group(2)
+        print(f"Invalid email format: {email}. Using default Telegram chat ID.")
+        return DEFAULT_CHAT_ID, 'telegram'
+
+    async def send_telegram(self, chat_id, message):
+        bot = telegram.Bot(TELEGRAM_API_KEY)
+        try:
+            await bot.send_message(chat_id=chat_id, text=message)
+            print(f"Message sent to Telegram chat ID: {chat_id}")
+        except telegram.error.BadRequest as e:
+            print(f"Failed to send message to Telegram chat ID: {chat_id}. Error: {str(e)}")
+            print("Attempting to send to default Telegram chat ID.")
+            await bot.send_message(chat_id=DEFAULT_CHAT_ID, text=message)
+
+    async def send_ntfy(self, topic, subject, body):
+        headers = {
+            "Authorization": f"Bearer {NTFY_TOKEN}",
+            "Title": subject,
+            "Content-Type": "text/plain",
+        }
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.post(f"{NTFY_URL}/{topic}", content=body, headers=headers)
+                response.raise_for_status()
+                print(f"Message sent to ntfy topic: {topic}")
+            except httpx.HTTPStatusError as e:
+                print(f"Failed to send message to ntfy topic: {topic}. Error: {str(e)}")
 
 async def start_smtp_server():
     controller = Controller(SMTPHandler(), hostname=SMTP_HOST, port=SMTP_PORT)