From d3bb82aedbcd8aa3d711ed72ab3f1e70ca2dc592 Mon Sep 17 00:00:00 2001
From: iwumingz <iwuming@outlook.com>
Date: Fri, 8 Apr 2022 14:53:17 +0800
Subject: [PATCH] Update version and add update plugin

---
 core/custom.py     |  2 +-
 plugins/update.py  | 33 +++++++++++++++++++++++++++++++++
 tools/constants.py | 13 +++++++++++++
 tools/helpers.py   |  4 ++--
 4 files changed, 49 insertions(+), 3 deletions(-)
 create mode 100644 plugins/update.py

diff --git a/core/custom.py b/core/custom.py
index 4218d9f..80b348a 100644
--- a/core/custom.py
+++ b/core/custom.py
@@ -4,7 +4,7 @@
 @File    :   custom.py
 @Time    :   2022/04/02 10:17:03
 @Author  :   Viperorz
-@Version :   1.0.0
+@Version :   1.0.5
 @License :   (C)Copyright 2021-2022
 @Desc    :   None
 """
diff --git a/plugins/update.py b/plugins/update.py
new file mode 100644
index 0000000..bc1dba8
--- /dev/null
+++ b/plugins/update.py
@@ -0,0 +1,33 @@
+import asyncio
+import sys
+
+from core import command
+from pyrogram import Client
+from pyrogram.types import Message
+from tools.constants import SYCGRAM, SYCGRAM_ERROR, SYCGRAM_INFO, SYCGRAM_WARNING, UPDATE_CMD
+from tools.helpers import basher
+
+
+@Client.on_message(command("restart"))
+async def restart(_: Client, msg: Message):
+    """重启容器"""
+    text = f"**{SYCGRAM_INFO}**\n> # `The {SYCGRAM} is restarting......`"
+    await msg.edit_text(text=text, parse_mode='md')
+    sys.exit()
+
+
+@Client.on_message(command("update"))
+async def update(_: Client, msg: Message):
+    """更新sycgram到主分支的最新版本"""
+    text = f"**{SYCGRAM_INFO}**\n> # `It's updating container to the latest version......`"
+    await msg.edit_text(text, parse_mode='md')
+    try:
+        _ = await basher(UPDATE_CMD, timeout=60)
+    except asyncio.exceptions.TimeoutError:
+        text = f"**{SYCGRAM_WARNING}**\n> # `Update Timeout！`"
+    except Exception as e:
+        text = f"**{SYCGRAM_ERROR}**\n> # `{e}`"
+    else:
+        text = f"**{SYCGRAM_INFO}**\n> # `Your {SYCGRAM} version is the latest.`"
+    finally:
+        await msg.edit_text(text, parse_mode='md')
diff --git a/tools/constants.py b/tools/constants.py
index 8c6b47e..c6dbff0 100644
--- a/tools/constants.py
+++ b/tools/constants.py
@@ -1,5 +1,18 @@
 from typing import Dict, List
 
+# Some params of sycgram
+SYCGRAM: str = "sycgram"
+SYCGRAM_INFO: str = f"{SYCGRAM.title()} | INFO"
+SYCGRAM_ERROR: str = f"{SYCGRAM.title()} | ERROR"
+SYCGRAM_WARNING: str = f"{SYCGRAM.title()} | WARNING"
+UPDATE_CMD: str = f"""
+docker run --rm \
+    -v /var/run/docker.sock:/var/run/docker.sock \
+    containrrr/watchtower \
+    --cleanup \
+    --run-once \
+    {SYCGRAM}
+"""
 
 # ------------- rate --------------
 RATE_API: str = 'https://cdn.jsdelivr.net/gh/fawazahmed0/currency-api@1/latest/currencies'
diff --git a/tools/helpers.py b/tools/helpers.py
index 28fde9f..7d74e9a 100644
--- a/tools/helpers.py
+++ b/tools/helpers.py
@@ -156,8 +156,8 @@ async def execute(command: str) -> Dict[str, Any]:
         return {'output': '', 'error': str(e)}
     else:
         return {
-            'output': stdout.decode('utf-8', 'ignore'),
-            'error': stderr.decode('utf-8', 'ignore')
+            'output': stdout.decode('utf-8', 'ignore').strip(),
+            'error': stderr.decode('utf-8', 'ignore').strip()
         }