From 74498087b9a605b54e708f64f8c8570131d7efdd Mon Sep 17 00:00:00 2001
From: omg-xtao <100690902+omg-xtao@users.noreply.github.com>
Date: Sat, 20 Jan 2024 23:26:56 +0800
Subject: [PATCH] :sparkles: Support count recv and send times

---
 gram_core                |  2 +-
 plugins/system/status.py | 41 +++++++++++++++++++++++++++++++++++++---
 2 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/gram_core b/gram_core
index ba02f9cc..909c81b1 160000
--- a/gram_core
+++ b/gram_core
@@ -1 +1 @@
-Subproject commit ba02f9cc75b67189ff9c55a17da48f44a0a54531
+Subproject commit 909c81b14035f97dae6fa0b7ea474576ecc25fd7
diff --git a/plugins/system/status.py b/plugins/system/status.py
index 1a117a05..bc295a3a 100644
--- a/plugins/system/status.py
+++ b/plugins/system/status.py
@@ -2,10 +2,11 @@ import asyncio
 import os
 from platform import python_version
 from time import time
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional, Union
 
 import psutil
-from telegram import __version__
+from telegram import __version__, Update
+from telegram.ext import BaseHandler
 
 from git import Repo
 from git.exc import GitCommandError, InvalidGitRepositoryError, NoSuchPathError
@@ -14,14 +15,47 @@ from core.plugin import Plugin, handler
 from utils.log import logger
 
 if TYPE_CHECKING:
-    from telegram import Update
     from telegram.ext import ContextTypes
 
 
+class StatisticsHandler(BaseHandler):
+    def __init__(self, plugin: "Status"):
+        self._plugin = plugin
+        super().__init__(self.recv_callback)
+
+    def check_update(self, update: object) -> Optional[Union[bool, object]]:
+        if isinstance(update, Update):
+            return True
+        return False
+
+    async def recv_callback(self, _: "Update", __: "ContextTypes.DEFAULT_TYPE"):
+        self._plugin.recv_num += 1
+
+    async def send_callback(self, endpoint: str, _, __):
+        if not endpoint:
+            return
+        if isinstance(endpoint, str) and endpoint.startswith("send"):
+            self._plugin.send_num += 1
+
+
 class Status(Plugin):
     def __init__(self):
         self.pid = os.getpid()
         self.time_form = "%m/%d %H:%M"
+        self.type_handler = None
+        self.recv_num = 0
+        self.send_num = 0
+
+    async def initialize(self) -> None:
+        self.type_handler = StatisticsHandler(self)
+        self.application.telegram.add_handler(self.type_handler, group=-10)
+
+        @self.application.on_called_api
+        async def call(endpoint: str, _, __):
+            await self.type_handler.send_callback(endpoint, _, __)
+
+    async def shutdown(self) -> None:
+        self.application.telegram.remove_handler(self.type_handler, group=-10)
 
     @staticmethod
     def get_git_hash() -> str:
@@ -62,6 +96,7 @@ class Status(Plugin):
             f"CPU使用率: `{cpu_percent}%/{process_cpu_use}%` \n"
             f"当前使用的内存: `{memory_text}` \n"
             f"运行时间: `{self.get_bot_uptime(start_time)}` \n"
+            f"收发消息: ⬇️ {self.recv_num} ⬆️ {self.send_num} \n"
         )
         await message.reply_markdown_v2(text)