Add notify to slack / telegram

Author: Miracle05

.gitignore

@@ -102,4 +102,7 @@ ENV/
 
 # docker/helm
 docker/certificates/**/*
-rules.py
+rules.py
+
+# config file
+artifactory-cleanup.yaml

artifactory_cleanup/__init__.py

@@ -7,4 +7,4 @@ def register(rule):
     registry.register(rule)
 
 
-__version__ = "1.0.18"
+__version__ = "1.1.0"

artifactory_cleanup/cli.py

@@ -15,11 +15,12 @@
 )
 from artifactory_cleanup.base_url_session import BaseUrlSession
 from artifactory_cleanup.context_managers import get_context_managers
-from artifactory_cleanup.errors import InvalidConfigError
+from artifactory_cleanup.errors import InvalidConfigError, NotificationError
 from artifactory_cleanup.loaders import (
     PythonLoader,
     YamlConfigLoader,
 )
+from artifactory_cleanup.notifiers import ReportNotifier, SlackConfig, TelegramConfig
 
 requests.packages.urllib3.disable_warnings()
 
@@ -102,6 +103,44 @@ class ArtifactoryCleanupCLI(cli.Application):
         default=False,
     )
 
+    _slack_token = cli.SwitchAttr(
+        "--slack-token",
+        help="Slack bot token for report upload",
+        mandatory=False,
+        envname="ARTIFACTORY_CLEANUP_SLACK_TOKEN",
+        requires=["--slack-channel-id", "--output"],
+    )
+
+    _slack_channel_id = cli.SwitchAttr(
+        "--slack-channel-id",
+        help="Slack channel id for report upload",
+        mandatory=False,
+        envname="ARTIFACTORY_CLEANUP_SLACK_CHANNEL_ID",
+        requires=["--slack-token", "--output"],
+    )
+
+    _telegram_token = cli.SwitchAttr(
+        "--telegram-token",
+        help="Telegram bot token for report upload",
+        mandatory=False,
+        envname="ARTIFACTORY_CLEANUP_TELEGRAM_TOKEN",
+        requires=["--telegram-chat-id", "--output"],
+    )
+
+    _telegram_chat_id = cli.SwitchAttr(
+        "--telegram-chat-id",
+        help="Telegram chat id for report upload",
+        mandatory=False,
+        envname="ARTIFACTORY_CLEANUP_TELEGRAM_CHAT_ID",
+        requires=["--telegram-token", "--output"],
+    )
+
+    _notify_comment = cli.SwitchAttr(
+        "--notify-comment",
+        help="Comment attached to uploaded report",
+        mandatory=False,
+    )
+
     @property
     def VERSION(self):
         # To prevent circular imports
@@ -153,6 +192,36 @@ def _create_output_file(self, result, filename, format):
         with open(filename, "w", encoding="utf-8") as file:
             file.write(text)
 
+    def _has_notifiers(self) -> bool:
+        return any(
+            [
+                self._slack_token,
+                self._slack_channel_id,
+                self._telegram_token,
+                self._telegram_chat_id,
+            ]
+        )
+
+    def _build_notifier(self) -> ReportNotifier:
+        slack = None
+        if self._slack_token and self._slack_channel_id:
+            slack = SlackConfig(
+                token=self._slack_token, channel_id=self._slack_channel_id
+            )
+        telegram = None
+        if self._telegram_token and self._telegram_chat_id:
+            telegram = TelegramConfig(
+                token=self._telegram_token, chat_id=self._telegram_chat_id
+            )
+        return ReportNotifier(slack=slack, telegram=telegram)
+
+    def _apply_notifier_config(self, config):
+        self._slack_token = self._slack_token or config.get("slack_token", "")
+        self._slack_channel_id = self._slack_channel_id or config.get("slack_channel_id", "")
+        self._telegram_token = self._telegram_token or config.get("telegram_token", "")
+        self._telegram_chat_id = self._telegram_chat_id or config.get("telegram_chat_id", "")
+        self._notify_comment = self._notify_comment or config.get("notify_comment", "")
+
     def main(self):
         today = self._get_today()
         if self._load_rules:
@@ -166,6 +235,8 @@ def main(self):
             print(str(err), file=sys.stderr)
             sys.exit(1)
 
+        self._apply_notifier_config(loader.get_notifications())
+
         server, user, password, apikey = loader.get_connection()
         session = BaseUrlSession(server)
         if apikey:
@@ -218,6 +289,13 @@ def main(self):
 
         if self._output_file:
             self._create_output_file(result, self._output_file, self._output_format)
+            if self._has_notifiers():
+                try:
+                    notifier = self._build_notifier()
+                    notifier.send_file(self._output_file, comment=self._notify_comment)
+                except NotificationError as err:
+                    print(str(err), file=sys.stderr)
+                    sys.exit(1)
 
 
 if __name__ == "__main__":

artifactory_cleanup/errors.py

@@ -4,3 +4,7 @@ class ArtifactoryCleanupException(Exception):
 
 class InvalidConfigError(ArtifactoryCleanupException):
     pass
+
+
+class NotificationError(ArtifactoryCleanupException):
+    pass

artifactory_cleanup/loaders.py

@@ -94,11 +94,29 @@ def get_root_schema(self, rules):
         config_schema = cfgv.Map(
             "Config",
             None,
-            cfgv.NoAdditionalKeys(["server", "user", "password", "policies", "apikey"]),
+            cfgv.NoAdditionalKeys(
+                [
+                    "server",
+                    "user",
+                    "password",
+                    "policies",
+                    "apikey",
+                    "slack_token",
+                    "slack_channel_id",
+                    "telegram_token",
+                    "telegram_chat_id",
+                    "notify_comment",
+                ]
+            ),
             cfgv.Required("server", cfgv.check_string),
             # User and password required, if apikey missing
             cfgv.Conditional("user", cfgv.check_string, "apikey", cfgv.MISSING, False),
             cfgv.Conditional("password", cfgv.check_string, "apikey", cfgv.MISSING, False),
+            cfgv.Optional("slack_token", cfgv.check_string, ""),
+            cfgv.Optional("slack_channel_id", cfgv.check_string, ""),
+            cfgv.Optional("telegram_token", cfgv.check_string, ""),
+            cfgv.Optional("telegram_chat_id", cfgv.check_string, ""),
+            cfgv.Optional("notify_comment", cfgv.check_string, ""),
             cfgv.RequiredRecurse("policies", cfgv.Array(policy_schema)),
         )
 
@@ -198,6 +216,21 @@ def get_connection(self) -> Tuple[str, str, str, str]:
         apikey = os.path.expandvars(apikey)
         return server, user, password, apikey
 
+    def get_notifications(self) -> Dict[str, str]:
+        config = self.load(self.filepath)
+        notifications = config.get("artifactory-cleanup", {})
+        return {
+            "slack_token": os.path.expandvars(notifications.get("slack_token", "")),
+            "slack_channel_id": os.path.expandvars(
+                notifications.get("slack_channel_id", "")
+            ),
+            "telegram_token": os.path.expandvars(notifications.get("telegram_token", "")),
+            "telegram_chat_id": os.path.expandvars(
+                notifications.get("telegram_chat_id", "")
+            ),
+            "notify_comment": os.path.expandvars(notifications.get("notify_comment", "")),
+        }
+
 
 class PythonLoader:
     """

artifactory_cleanup/notifiers.py

@@ -0,0 +1,129 @@
+import os
+from dataclasses import dataclass
+from typing import Optional, List
+
+import requests
+
+from artifactory_cleanup.errors import NotificationError
+
+
+@dataclass
+class SlackConfig:
+    token: str
+    channel_id: str
+
+
+@dataclass
+class TelegramConfig:
+    token: str
+    chat_id: str
+
+
+class SlackNotifier:
+    def __init__(self, config: SlackConfig):
+        self._config = config
+
+    def _headers(self):
+        return {"Authorization": f"Bearer {self._config.token}"}
+
+    def send_file(self, filename: str, comment: Optional[str] = None):
+        if not os.path.exists(filename):
+            raise NotificationError(f"Report file not found: {filename}")
+
+        file_size = os.path.getsize(filename)
+        file_basename = os.path.basename(filename)
+        metadata_response = requests.post(
+            "https://slack.com/api/files.getUploadURLExternal",
+            headers=self._headers(),
+            data={
+                "filename": file_basename,
+                "length": file_size,
+            },
+            timeout=30,
+        )
+        metadata_response.raise_for_status()
+        metadata = metadata_response.json()
+        if not metadata.get("ok"):
+            raise NotificationError(
+                f"Slack upload URL request failed: {metadata.get('error', 'unknown error')}"
+            )
+
+        upload_url = metadata["upload_url"]
+        file_id = metadata["file_id"]
+
+        with open(filename, "rb") as report_file:
+            upload_response = requests.post(
+                upload_url,
+                data=report_file,
+                headers={"Content-Type": "application/octet-stream"},
+                timeout=60,
+            )
+            upload_response.raise_for_status()
+
+        complete_payload = {
+            "files": [{"id": file_id, "title": file_basename}],
+            "channel_id": self._config.channel_id,
+        }
+        if comment:
+            complete_payload["initial_comment"] = comment
+
+        complete_response = requests.post(
+            "https://slack.com/api/files.completeUploadExternal",
+            headers={
+                **self._headers(),
+                "Content-Type": "application/json; charset=utf-8",
+            },
+            json=complete_payload,
+            timeout=30,
+        )
+        complete_response.raise_for_status()
+        complete = complete_response.json()
+        if not complete.get("ok"):
+            raise NotificationError(
+                f"Slack upload finalize failed: {complete.get('error', 'unknown error')}"
+            )
+
+
+class TelegramNotifier:
+    def __init__(self, config: TelegramConfig):
+        self._config = config
+
+    def send_file(self, filename: str, comment: Optional[str] = None):
+        if not os.path.exists(filename):
+            raise NotificationError(f"Report file not found: {filename}")
+
+        url = f"https://api.telegram.org/bot{self._config.token}/sendDocument"
+        data = {"chat_id": self._config.chat_id}
+        if comment:
+            data["caption"] = comment
+
+        with open(filename, "rb") as report_file:
+            response = requests.post(
+                url,
+                data=data,
+                files={"document": report_file},
+                timeout=60,
+            )
+        response.raise_for_status()
+        payload = response.json()
+        if not payload.get("ok"):
+            raise NotificationError(
+                f"Telegram send failed: {payload.get('description', 'unknown error')}"
+            )
+
+
+class ReportNotifier:
+    def __init__(
+        self,
+        slack: Optional[SlackConfig] = None,
+        telegram: Optional[TelegramConfig] = None,
+    ):
+        self._notifiers: List = []
+        if slack:
+            self._notifiers.append(SlackNotifier(slack))
+        if telegram:
+            self._notifiers.append(TelegramNotifier(telegram))
+
+    def send_file(self, filename: str, comment: Optional[str] = None):
+        for notifier in self._notifiers:
+            notifier.send_file(filename=filename, comment=comment)

setup.py

@@ -12,7 +12,7 @@
 
 setup(
     name="artifactory-cleanup",
-    version="1.0.18",
+    version="1.1.0",
     description="Rules and cleanup policies for Artifactory",
     long_description=long_description,
     long_description_content_type="text/markdown",