Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[FEATURE] Discord Notifications Fixes #186 #193

Merged
merged 3 commits into from
Feb 24, 2024
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
Added webhook-based discord notifications
  • Loading branch information
mandarons committed Feb 24, 2024
commit 0ff9abb90ce1604812f3bcd980065ae6626ed5ec
3 changes: 3 additions & 0 deletions config.yaml
Original file line number Diff line number Diff line change
@@ -11,6 +11,9 @@ app:
retry_login_interval: 600
# Drive destination
root: "icloud"
discord:
# webhook_url: <your server webhook URL here>
# username: icloud-docker #or any other name you prefer
telegram:
# bot_token: <your Telegram bot token>
# chat_id: <your Telegram user or chat ID>
28 changes: 28 additions & 0 deletions src/config_parser.py
Original file line number Diff line number Diff line change
@@ -418,3 +418,31 @@ def get_telegram_chat_id(config):
else:
chat_id = get_config_value(config=config, config_path=config_path)
return chat_id


# Get discord webhook_url
def get_discord_webhook_url(config):
"""Return discord webhook_url from config."""
webhook_url = None
config_path = ["app", "discord", "webhook_url"]
if not traverse_config_path(config=config, config_path=config_path):
LOGGER.warning(
f"Warning: webhook_url is not found in {config_path_to_string(config_path)}."
)
else:
webhook_url = get_config_value(config=config, config_path=config_path)
return webhook_url


# Get discord username
def get_discord_username(config):
"""Return discord username from config."""
username = None
config_path = ["app", "discord", "username"]
if not traverse_config_path(config=config, config_path=config_path):
LOGGER.warning(
f"Warning: username is not found in {config_path_to_string(config_path)}."
)
else:
username = get_config_value(config=config, config_path=config_path)
return username
49 changes: 42 additions & 7 deletions src/notify.py
Original file line number Diff line number Diff line change
@@ -7,6 +7,11 @@
from src import LOGGER, config_parser
from src.email_message import EmailMessage as Message

MESSAGE_BODY = """Two-step authentication for iCloud Drive, Photos (Docker) is required.
Please login to your server and authenticate. Please run -
`docker exec -it icloud /bin/sh -c "icloud --username=<icloud-username>
--session-directory=/app/session_data"`."""


def notify_telegram(config, last_send=None, dry_run=False):
"""Send telegram notification."""
@@ -24,10 +29,7 @@ def notify_telegram(config, last_send=None, dry_run=False):
if not post_message_to_telegram(
bot_token,
chat_id,
"""Two-step authentication for iCloud Drive, Photos (Docker) is required.
Please login to your server and authenticate. Please run -
`docker exec -it icloud /bin/sh -c "icloud --username=<icloud-username>
--session-directory=/app/session_data"`.""",
MESSAGE_BODY,
):
sent_on = None
else:
@@ -44,16 +46,49 @@ def post_message_to_telegram(bot_token, chat_id, message):
response = requests.post(url, params=params, timeout=10)
if response.status_code == 200:
return True
# Log error message
LOGGER.error(f"Failed to send telegram notification. Response: {response.text}")
return False


def post_message_to_discord(webhook_url, username):
"""Post message to discord webhook."""
data = {"username": username, "content": MESSAGE_BODY}
response = requests.post(webhook_url, data=data, timeout=10)
if response.status_code == 204:
return True
# Log error message
LOGGER.error(f"Failed to send telegram notification. Response: {response.text}")
return False


def notify_discord(config, last_send=None, dry_run=False):
"""Send discord notification."""
sent_on = None
webhook_url = config_parser.get_discord_webhook_url(config=config)
username = config_parser.get_discord_username(config=config)

if last_send and last_send > datetime.datetime.now() - datetime.timedelta(hours=24):
LOGGER.info("Throttling discord to once a day")
sent_on = last_send
elif webhook_url and username:
sent_on = datetime.datetime.now()
if not dry_run:
# Post message to discord webhook using API
if not post_message_to_discord(webhook_url, username):
sent_on = None
else:
# Log error message
LOGGER.error(f"Failed to send telegram notification. Response: {response.text}")
return False
LOGGER.warning(
"Not sending 2FA notification because Discord is not configured."
)
return sent_on


def send(config, last_send=None, dry_run=False):
"""Send notifications."""
sent_on = None
notify_telegram(config=config, last_send=last_send, dry_run=dry_run)
notify_discord(config=config, last_send=last_send, dry_run=dry_run)
email = config_parser.get_smtp_email(config=config)
to_email = config_parser.get_smtp_to_email(config=config)
host = config_parser.get_smtp_host(config=config)
3 changes: 3 additions & 0 deletions tests/data/test_config.yaml
Original file line number Diff line number Diff line change
@@ -11,6 +11,9 @@ app:
retry_login_interval: 600
# Drive destination
root: "./icloud"
discord:
webhook_url: <server webhook>
username: icloud-docker
smtp:
# If you want to recieve email notifications about expired/missing 2FA credentials then uncomment
# email: sender@test.com
28 changes: 28 additions & 0 deletions tests/test_config_parser.py
Original file line number Diff line number Diff line change
@@ -489,3 +489,31 @@ def test_get_telegram_chat_id(self):
def test_get_telegram_chat_id_none_config(self):
"""None config."""
self.assertIsNone(config_parser.get_telegram_chat_id(config=None))

# write a test for discord webhook url
def test_get_discord_webhook_url(self):
"""Test for discord webhook url."""
config = read_config(config_path=tests.CONFIG_PATH)
config["app"]["discord"] = {"webhook_url": "webhook_url"}
self.assertEqual(
config["app"]["discord"]["webhook_url"],
config_parser.get_discord_webhook_url(config=config),
)

def test_get_discord_webhook_url_none_config(self):
"""None config."""
self.assertIsNone(config_parser.get_discord_webhook_url(config=None))

# write a test for discord username
def test_get_discord_username(self):
"""Test for discord username."""
config = read_config(config_path=tests.CONFIG_PATH)
config["app"]["discord"] = {"username": "username"}
self.assertEqual(
config["app"]["discord"]["username"],
config_parser.get_discord_username(config=config),
)

def test_get_discord_username_none_config(self):
"""None config."""
self.assertIsNone(config_parser.get_discord_username(config=None))
118 changes: 108 additions & 10 deletions tests/test_notify.py
Original file line number Diff line number Diff line change
@@ -5,7 +5,13 @@

from src import config_parser, notify
from src.email_message import EmailMessage as Message
from src.notify import notify_telegram, post_message_to_telegram
from src.notify import (
MESSAGE_BODY,
notify_discord,
notify_telegram,
post_message_to_discord,
post_message_to_telegram,
)


class TestNotify(unittest.TestCase):
@@ -139,10 +145,7 @@ def test_notify_telegram_success(self):
post_message_mock.assert_called_once_with(
config["app"]["telegram"]["bot_token"],
config["app"]["telegram"]["chat_id"],
"""Two-step authentication for iCloud Drive, Photos (Docker) is required.
Please login to your server and authenticate. Please run -
`docker exec -it icloud /bin/sh -c "icloud --username=<icloud-username>
--session-directory=/app/session_data"`.""",
MESSAGE_BODY,
)

def test_notify_telegram_fail(self):
@@ -161,18 +164,15 @@ def test_notify_telegram_fail(self):
post_message_mock.assert_called_once_with(
config["app"]["telegram"]["bot_token"],
config["app"]["telegram"]["chat_id"],
"""Two-step authentication for iCloud Drive, Photos (Docker) is required.
Please login to your server and authenticate. Please run -
`docker exec -it icloud /bin/sh -c "icloud --username=<icloud-username>
--session-directory=/app/session_data"`.""",
MESSAGE_BODY,
)

def test_notify_telegram_throttling(self):
"""Test for throttled notification."""
config = {
"telegram": {"bot_token": "your-bot-token", "chat_id": "your-chat-id"}
}
last_send = datetime.datetime.now() - datetime.timedelta(hours=24)
last_send = datetime.datetime.now() - datetime.timedelta(hours=2)
dry_run = False

with patch("src.notify.post_message_to_telegram") as post_message_mock:
@@ -232,3 +232,101 @@ def test_post_message_to_telegram_fail(self):
params={"chat_id": "chat_id", "text": "message"},
timeout=10,
)

def test_notify_discord_success(self):
"""Test for successful notification."""
config = {
"app": {"discord": {"webhook_url": "webhook-url", "username": "username"}}
}

with patch("src.notify.post_message_to_discord") as post_message_mock:
notify_discord(config, None, False)

# Verify that post_message_to_discord is called with the correct arguments
post_message_mock.assert_called_once_with(
config["app"]["discord"]["webhook_url"],
config["app"]["discord"]["username"],
)
self.assertEqual(post_message_mock.call_count, 1)

def test_notify_discord_fail(self):
"""Test for failed notification."""
config = {
"app": {"discord": {"webhook_url": "webhook-url", "username": "username"}}
}

with patch("src.notify.post_message_to_discord") as post_message_mock:
post_message_mock.return_value = False
notify_discord(config, None, False)

# Verify that post_message_to_discord is called with the correct arguments
post_message_mock.assert_called_once_with(
config["app"]["discord"]["webhook_url"],
config["app"]["discord"]["username"],
)

def test_notify_discord_throttling(self):
"""Test for throttled notification."""
config = {
"app": {"discord": {"webhook_url": "webhook-url", "username": "username"}}
}
last_send = datetime.datetime.now() - datetime.timedelta(hours=2)
dry_run = False

with patch("src.notify.post_message_to_discord") as post_message_mock:
notify_discord(config, last_send, dry_run)

# Verify that post_message_to_discord is not called when throttled
post_message_mock.assert_not_called()

def test_notify_discord_dry_run(self):
"""Test for dry run mode."""
config = {
"app": {"discord": {"webhook_url": "webhook-url", "username": "username"}}
}
last_send = datetime.datetime.now()
dry_run = True

with patch("src.notify.post_message_to_discord") as post_message_mock:
notify_discord(config, last_send, dry_run)

# Verify that post_message_to_discord is not called in dry run mode
post_message_mock.assert_not_called()

def test_notify_discord_no_config(self):
"""Test for missing discord configuration."""
config = {}
last_send = None
dry_run = False

with patch("src.notify.post_message_to_discord") as post_message_mock:
notify_discord(config, last_send, dry_run)

# Verify that post_message_to_discord is not called when discord configuration is missing
post_message_mock.assert_not_called()

def test_post_message_to_discord(self):
"""Test for successful post."""
with patch("requests.post") as post_mock:
post_mock.return_value.status_code = 200
post_message_to_discord("webhook_url", "username")

# Verify that post is called with the correct arguments
post_mock.assert_called_once_with(
"webhook_url",
data={"content": MESSAGE_BODY, "username": "username"},
timeout=10,
)

def test_post_message_to_discord_fail(self):
"""Test for failed post."""
with patch("requests.post") as post_mock:
post_mock.return_value.status_code = 400
post_message_to_discord("webhook_url", "username")

# Verify that post is called with the correct arguments
post_mock.assert_called_once_with(
"webhook_url",
data={"content": MESSAGE_BODY, "username": "username"},
timeout=10,
)
Loading