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

Slack alert throttle #15

Merged
merged 14 commits into from
Nov 18, 2024
3 changes: 3 additions & 0 deletions files/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
psycopg2-binary
neoformit marked this conversation as resolved.
Show resolved Hide resolved
requests
pyyaml
59 changes: 59 additions & 0 deletions files/test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
import unittest
neoformit marked this conversation as resolved.
Show resolved Hide resolved
from unittest.mock import patch
from datetime import datetime, timedelta
import pathlib
import tempfile
from walle import NotificationHistory

SLACK_NOTIFY_PERIOD_DAYS = 7


class TestNotificationHistory(unittest.TestCase):
def setUp(self):
self.temp_file = tempfile.NamedTemporaryFile(delete=False)
self.record = NotificationHistory(self.temp_file.name)

def tearDown(self):
pathlib.Path(self.temp_file.name).unlink(missing_ok=True)

def test_contains_new_entry(self):
jwd = "unique_id_1"
self.assertFalse(
self.record.contains(jwd), "New entry should initially return False"
)
self.assertTrue(self.record.contains(jwd), "Duplicate entry should return True")

def test_contains_existing_entry(self):
jwd = "existing_id"
with open(self.temp_file.name, "a") as f:
f.write(f"{datetime.now()}\t{jwd}\n")
self.assertTrue(self.record.contains(jwd), "Existing entry should return True")

@patch("walle.SLACK_NOTIFY_PERIOD_DAYS", new=SLACK_NOTIFY_PERIOD_DAYS)
def test_truncate_old_records(self):
old_jwd = "old_entry"
recent_jwd = "recent_entry"
old_date = datetime.now() - timedelta(days=SLACK_NOTIFY_PERIOD_DAYS + 1)
recent_date = datetime.now()

with open(self.temp_file.name, "a") as f:
f.write(f"{old_date.isoformat()}\t{old_jwd}\n")
f.write(f"{recent_date.isoformat()}\t{recent_jwd}\n")

self.record._truncate_records()
self.assertFalse(self.record.contains(old_jwd), "Old entry should be purged")
self.assertTrue(self.record.contains(recent_jwd), "Recent entry should remain")

def test_purge_invalid_records(self):
with open(self.temp_file.name, "w") as f:
f.write("invalid_date\tinvalid_path\n")

with patch("walle.logger.warning") as mock_warning:
self.record._read_records()
mock_warning.assert_called()

self.assertFalse(self.record._get_jwds(), "Invalid records should be purged")


if __name__ == "__main__":
unittest.main()
74 changes: 73 additions & 1 deletion files/walle.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@
import sys
import time
import zlib
from typing import Dict, List
from datetime import datetime, timedelta
from typing import Dict, List, Union

import galaxy_jwd
import requests
Expand All @@ -33,6 +34,9 @@
If you think your account was deleted due to an error, please contact
"""
ONLY_ONE_INSTANCE = "The other must be an instance of the Severity class"

# Number of days before repeating slack alert for the same JWD
SLACK_NOTIFY_PERIOD_DAYS = 7
SLACK_URL = "https://slack.com/api/chat.postMessage"

UserId = str
Expand All @@ -46,6 +50,9 @@
)
logger = logging.getLogger(__name__)
GXADMIN_PATH = os.getenv("GXADMIN_PATH", "/usr/local/bin/gxadmin")
NOTIFICATION_HISTORY_FILE = os.getenv(
"WALLE_NOTIFICATION_HISTORY_FILE", "/tmp/walle-notifications.txt"
)


def convert_arg_to_byte(mb: str) -> int:
Expand All @@ -56,6 +63,67 @@ def convert_arg_to_seconds(hours: str) -> float:
return float(hours) * 60 * 60


class NotificationHistory:
"""Record of Slack notifications to avoid spamming users."""

def __init__(self, record_file: str) -> None:
self.record_file = pathlib.Path(record_file)
if not self.record_file.exists():
self.record_file.touch()
self._truncate_records()

def _get_jwds(self) -> List[str]:
return [line[1] for line in self._read_records()]

def _read_records(self) -> List[List[str]]:
with open(self.record_file, "r") as f:
records = [
line.strip().split("\t") for line in f.readlines() if line.strip()
]
return self._validate(records)

def _validate(self, records: List[List[str]]) -> List[List[str]]:
neoformit marked this conversation as resolved.
Show resolved Hide resolved
try:
for datestr, path in records:
if not isinstance(datestr, str) and isinstance(path, str):
raise ValueError
datetime.fromisoformat(datestr)
except ValueError:
logger.warning(
f"Invalid records found in {self.record_file}. The"
" file will be purged. This may result in duplicate Slack"
" notifications."
)
self._purge_records()
return []
return records

def _write_jwd(self, jwd: str) -> None:
with open(self.record_file, "a") as f:
f.write(f"{datetime.now()}\t{jwd}\n")

def _purge_records(self) -> None:
self.record_file.unlink()
self.record_file.touch()

def _truncate_records(self) -> None:
"""Truncate older records."""
records = self._read_records()
with open(self.record_file, "w") as f:
for datestr, jwd_path in records:
if datetime.fromisoformat(datestr) > datetime.now() - timedelta(
days=SLACK_NOTIFY_PERIOD_DAYS
):
f.write(f"{datestr}\t{jwd_path}\n")

def contains(self, jwd: Union[pathlib.Path, str]) -> bool:
jwd = str(jwd)
exists = jwd in self._get_jwds()
if not exists:
self._write_jwd(jwd)
return exists


class Severity:
def __init__(self, number: int, name: str):
self.value = number
Expand Down Expand Up @@ -87,6 +155,7 @@ def __ge__(self, other) -> bool:


VALID_SEVERITIES = (Severity(0, "LOW"), Severity(1, "MEDIUM"), Severity(2, "HIGH"))
notification_history = NotificationHistory(NOTIFICATION_HISTORY_FILE)


def convert_str_to_severity(test_level: str) -> Severity:
Expand Down Expand Up @@ -406,6 +475,9 @@ def report_matching_malware(self):
)

def post_slack_alert(self):
if notification_history.contains(self.job.jwd):
logger.debug("Skipping Slack notification - already posted for this JWD")
return
msg = f"""
:rotating_light: WALLE: *Malware detected* :rotating_light:

Expand Down
Loading