From 12c124863e9920bc10c2093f4f1ff027860a3b0a Mon Sep 17 00:00:00 2001
From: sergey <sergey@sm8255082.ru>
Date: Sun, 13 Apr 2025 15:30:37 +0300
Subject: [PATCH] add fail check and edit message

---
 config/.env-template       |  1 +
 config/config.py           |  2 ++
 docker/docker-compose.yaml |  2 +-
 pyproject.toml             |  2 +-
 telegram/__init__.py       |  2 ++
 telegram/message.py        | 32 ++++++++++++++++++++++++++++++++
 zabbix/zabbix_api.py       |  7 +++++++
 7 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/config/.env-template b/config/.env-template
index 8b59671..cf0bd90 100755
--- a/config/.env-template
+++ b/config/.env-template
@@ -5,6 +5,7 @@ CFG__ZABBIX__URL=https://example.com
 CFG__ZABBIX__TOKEN=string
 CFG__ZABBIX__MIN_SEVERITY=0
 CFG__ZABBIX__UPD_INTERVAL=60
+CFG__ZABBIX__ALERT_FAIL_COUNT=3
 
 CFG__TGBOT__TOKEN=string
 CFG__TGBOT__CHAT_ID=00000000
diff --git a/config/config.py b/config/config.py
index 9c0c501..d41d4e4 100755
--- a/config/config.py
+++ b/config/config.py
@@ -49,6 +49,8 @@ class ZabbixConfig(BaseModel):
     token: str
     min_severity: int
     upd_interval: int
+    alert_fail_count: int
+    current_fail: int = 0
 
 
 class Settings(BaseSettings):
diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml
index 369c50e..b9a6426 100644
--- a/docker/docker-compose.yaml
+++ b/docker/docker-compose.yaml
@@ -10,7 +10,7 @@ services:
     command: [redis-server, --protected-mode yes, --port 6379, --requirepass, P@ssw0rd!]
 
   tg-bot:
-    image: git.sm8255082.ru/osnova/zbx-tg-bot:2.0.0
+    image: git.sm8255082.ru/osnova/zbx-tg-bot:2.0.1
     restart: always
     depends_on:
       - redis
diff --git a/pyproject.toml b/pyproject.toml
index 0c7b748..069f818 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "zbx-tg-bot"
-version = "2.0.0"
+version = "2.0.1"
 description = "telegram bot for telegram-zabbix dashboard"
 requires-python = ">=3.13"
 dependencies = [
diff --git a/telegram/__init__.py b/telegram/__init__.py
index 9a89e27..3427200 100644
--- a/telegram/__init__.py
+++ b/telegram/__init__.py
@@ -1,6 +1,7 @@
 from .message import (
     send_message,
     del_message,
+    edit_message,
 )
 from .bot import start_bot
 
@@ -8,4 +9,5 @@ __all__ = [
     "send_message",
     "del_message",
     "start_bot",
+    "edit_message",
 ]
diff --git a/telegram/message.py b/telegram/message.py
index 2a3922f..c6c0650 100644
--- a/telegram/message.py
+++ b/telegram/message.py
@@ -45,6 +45,38 @@ async def send_message(message: str, event_id: int) -> dict:
             return {"status": e}
 
 
+async def edit_message(
+    message_id: int,
+    new_message: str,
+) -> dict:
+    url = f"https://api.telegram.org/bot{conf.tgbot.token}/editMessageText"
+    async with aiohttp.ClientSession() as session:
+        try:
+            async with session.post(
+                url,
+                json={
+                    "chat_id": conf.tgbot.chat_id,
+                    "message_id": message_id,
+                    "text": new_message,
+                },
+            ) as response:
+                if response.status == 200:
+                    log.info(f"Message ID {message_id} edited")
+                    return {
+                        "status": response.status,
+                    }
+                else:
+                    log.warning(
+                        f"Message ID {message_id} NOT edited. Response status: {response.status}"
+                    )
+                    return {
+                        "status": response.status,
+                    }
+        except Exception as e:
+            log.warning(f"Exception: {e}")
+            return {"status": e}
+
+
 async def del_message(
     message_id: int,
 ) -> dict:
diff --git a/zabbix/zabbix_api.py b/zabbix/zabbix_api.py
index 44c6b2a..4f345a9 100644
--- a/zabbix/zabbix_api.py
+++ b/zabbix/zabbix_api.py
@@ -40,9 +40,16 @@ def get_active_problems() -> dict:
             event["host"] = event.pop("hosts", None)[0]["host"]
             events_dict[event["eventid"]] = event
             events_dict["event_ids"].append(event["eventid"])
+        if conf.zabbix.current_fail:
+            conf.zabbix.current_fail = 0
         return events_dict
     except:
         log.warning("Get event from zabbix error")
+        if conf.zabbix.current_fail == conf.zabbix.alert_fail_count:
+            log.warning(f"Zabbix connection error. Count {conf.zabbix.current_fail}")
+        else:
+            conf.zabbix.current_fail += 1
+            log.warning(f"Zabbix connection error. Count {conf.zabbix.current_fail}")
 
 
 def event_acknowledge(event_id: int, mute_time: int, muted_by: str):