diff options
| -rw-r--r-- | .github/workflows/lint.yml | 2 | ||||
| -rw-r--r-- | ATRI/exceptions.py | 10 | ||||
| -rw-r--r-- | ATRI/plugins/console/__init__.py | 202 | ||||
| -rw-r--r-- | ATRI/plugins/console/data_source.py | 58 | ||||
| -rw-r--r-- | ATRI/plugins/console/driver/__init__.py | 56 | ||||
| -rw-r--r-- | ATRI/plugins/console/driver/api.py | 193 | ||||
| -rw-r--r-- | ATRI/plugins/console/driver/view.py | 101 | ||||
| -rw-r--r-- | ATRI/plugins/console/listener.py | 49 | ||||
| -rw-r--r-- | ATRI/plugins/console/models.py | 40 | ||||
| -rw-r--r-- | ATRI/plugins/kimo/__init__.py | 3 | ||||
| -rw-r--r-- | ATRI/plugins/kimo/data_source.py | 6 | ||||
| -rw-r--r-- | ATRI/plugins/polaroid/data_source.py | 6 | ||||
| -rw-r--r-- | ATRI/plugins/setu/tf_dealer.py | 10 | ||||
| -rw-r--r-- | ATRI/plugins/status/__init__.py | 20 | ||||
| -rw-r--r-- | ATRI/plugins/status/data_source.py | 47 | ||||
| -rw-r--r-- | ATRI/service.py | 18 | 
16 files changed, 765 insertions, 56 deletions
| diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index c8eb30f..32499dd 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -39,7 +39,7 @@ jobs:            python -m pip install --upgrade pip            pip install flake8        - name: Run Linters -        uses: wearerequired/lint-action@v1 +        uses: wearerequired/lint-action@v2.0.0          with:            auto_fix: true            black: true diff --git a/ATRI/exceptions.py b/ATRI/exceptions.py index 8ce6a85..4d52ff3 100644 --- a/ATRI/exceptions.py +++ b/ATRI/exceptions.py @@ -62,12 +62,12 @@ class InvalidConfigured(BaseBotException):      prompt = "无效配置" -class WriteError(BaseBotException): +class WriteFileError(BaseBotException):      prompt = "写入错误" -class LoadingError(BaseBotException): -    prompt = "加载错误" +class ReadFileError(BaseBotException): +    prompt = "读取文件失败"  class RequestError(BaseBotException): @@ -78,10 +78,6 @@ class GetStatusError(BaseBotException):      prompt = "获取状态失败" -class ReadFileError(BaseBotException): -    prompt = "读取文件失败" - -  class FormatError(BaseBotException):      prompt = "格式错误" diff --git a/ATRI/plugins/console/__init__.py b/ATRI/plugins/console/__init__.py new file mode 100644 index 0000000..c6fd847 --- /dev/null +++ b/ATRI/plugins/console/__init__.py @@ -0,0 +1,202 @@ +import json + +from nonebot.params import ArgPlainText +from nonebot.adapters.onebot.v11 import PrivateMessageEvent, GroupMessageEvent + +from ATRI.config import BotSelfConfig +from ATRI.exceptions import WriteFileError, ReadFileError +from .data_source import Console, CONSOLE_DIR +from .models import AuthData +from .driver import init + +init() + + +gen_console_key = Console().cmd_as_group("auth", "获取进入网页后台的凭证") + + +@gen_console_key.got("is_pub_n", "咱的运行环境是否有公网(y/n)") +async def _(event: PrivateMessageEvent, is_pub_n: str = ArgPlainText("is_pub_n")): +    if is_pub_n != "y": +        ip = str(await Console().get_host_ip(False)) +        await gen_console_key.send("没有公网吗...嗯知道了") +    else: +        ip = str(await Console().get_host_ip(True)) + +    p = BotSelfConfig.port +    rs = Console().get_random_str(20) + +    df = CONSOLE_DIR / "data.json" +    try: +        if not df.is_file(): +            with open(df, "w", encoding="utf-8") as w: +                w.write(json.dumps({})) + +        d = json.loads(df.read_bytes()) + +        ca = d.get("data", None) +        if ca: +            # 此处原本想用 matcher.finish 但这是在 try 里啊! +            await gen_console_key.send("咱已经告诉你了嗷!啊!忘了.../gauth 获取吧") +            return + +        d["data"] = AuthData(ip=ip, port=str(p), token=rs).dict() + +        with open(df, "w", encoding="utf-8") as w: +            w.write(json.dumps(d)) +    except WriteFileError: +        msg = f""" +        哦吼!写入文件失败了...还请自行记下哦... +        IP: {ip} +        PORT: {p} +        TOKEN: {rs} +        一定要保管好哦!切勿告诉他人哦! +        """.strip() +        await gen_console_key.send(msg) + +        raise WriteFileError("Writing file: " + str(df) + " failed!") + +    msg = f""" +    该信息已保存!可通过 /gauth 获取~ +    IP: {ip} +    PORT: {p} +    TOKEN: {rs} +    一定要保管好哦!切勿告诉他人哦! +    """.strip() +    await gen_console_key.finish(msg) + + +@gen_console_key.handle() +async def _(event: GroupMessageEvent): +    await gen_console_key.finish("请私戳咱获取(") + + +load_console_key = Console().cmd_as_group("load", "获取已生成的后台凭证") + + +@load_console_key.handle() +async def _(event: PrivateMessageEvent): +    df = CONSOLE_DIR / "data.json" +    if not df.is_file(): +        await load_console_key.finish("你还没有问咱索要奥!/auth 以获取") + +    try: +        d = json.loads(df.read_bytes()) +    except ReadFileError: +        await load_console_key.send("获取数据失败了...请自行打开文件查看吧:\n" + str(df)) +        raise ReadFileError("Reading file: " + str(df) + " failed!") + +    data = d["data"] +    msg = f""" +    诶嘿嘿嘿——凭证信息来咯! +    IP: {data['ip']} +    PORT: {data['port']} +    TOKEN: {data['token']} +    切记!不要告诉他人!! +    """.strip() +    await load_console_key.finish(msg) + + +@load_console_key.handle() +async def _(event: GroupMessageEvent): +    await load_console_key.finish("请私戳咱获取(") + + +del_console_key = Console().cmd_as_group("del", "销毁进入网页后台的凭证") + + +@del_console_key.got("is_sure_d", "...你确定吗(y/n)") +async def _(is_sure: str = ArgPlainText("is_sure_d")): +    if is_sure != "y": +        await del_console_key.finish("反悔了呢...") + +    df = CONSOLE_DIR / "data.json" +    if not df.is_file(): +        await del_console_key.finish("你还没向咱索取凭证呢.../auth 以获取") + +    try: +        data: dict = json.loads(df.read_bytes()) + +        del data["data"] + +        with open(df, "w", encoding="utf-8") as w: +            w.write(json.dumps(data)) +    except WriteFileError: +        await del_console_key.send("销毁失败了...请至此处自行删除文件:\n" + str(df)) +        raise WriteFileError("Writing / Reading file: " + str(df) + " failed!") + +    await del_console_key.finish("销毁成功!如需再次获取: /auth") + + +res_console_key = Console().cmd_as_group("reauth", "重置进入网页后台的凭证") + + +@res_console_key.got("is_sure_r", "...你确定吗(y/n)") +async def _(is_sure: str = ArgPlainText("is_sure_r")): +    if is_sure != "y": +        await res_console_key.finish("反悔了呢...") + +    df = CONSOLE_DIR / "data.json" +    if not df.is_file(): +        await del_console_key.finish("你还没向咱索取凭证呢.../auth 以获取") + +    try: +        data: dict = json.loads(df.read_bytes()) + +        del data["data"] + +        with open(df, "w", encoding="utf-8") as w: +            w.write(json.dumps(data)) +    except WriteFileError: +        await del_console_key.send("销毁失败了...请至此处自行删除文件:\n" + str(df)) +        raise WriteFileError("Writing / Reading file: " + str(df) + " failed!") + + +@res_console_key.got("is_pub_r_n", "咱的运行环境是否有公网(y/n)") +async def _(event: PrivateMessageEvent, is_pub_n: str = ArgPlainText("is_pub_n")): +    if is_pub_n != "y": +        ip = str(await Console().get_host_ip(False)) +        await res_console_key.send("没有公网吗...嗯知道了") +    else: +        ip = str(await Console().get_host_ip(True)) + +    p = BotSelfConfig.port +    rs = Console().get_random_str(20) + +    df = CONSOLE_DIR / "data.json" +    try: +        if not df.is_file(): +            with open(df, "w", encoding="utf-8") as w: +                w.write(json.dumps({})) + +        d = json.loads(df.read_bytes()) + +        ca = d.get("data", None) +        if ca: +            await res_console_key.send("咱已经告诉你了嗷!啊!忘了.../gauth 获取吧") +            return + +        d["data"] = AuthData(ip=ip, port=str(p), token=rs).dict() + +        with open(df, "w", encoding="utf-8") as w: +            w.write(json.dumps(d)) +    except WriteFileError: +        msg = f""" +        哦吼!写入文件失败了...还请自行记下哦... +        IP: {ip} +        PORT: {p} +        TOKEN: {rs} +        一定要保管好哦!切勿告诉他人哦! +        """.strip() +        await res_console_key.send(msg) + +        raise WriteFileError("Writing file: " + str(df) + " failed!") + +    msg = f""" +    该信息已保存!可通过 /gauth 获取~ +    IP: {ip} +    PORT: {p} +    TOKEN: {rs} +    一定要保管好哦!切勿告诉他人哦! +    """.strip() +    await res_console_key.finish(msg) diff --git a/ATRI/plugins/console/data_source.py b/ATRI/plugins/console/data_source.py new file mode 100644 index 0000000..86a3813 --- /dev/null +++ b/ATRI/plugins/console/data_source.py @@ -0,0 +1,58 @@ +import json +import socket +import string +from random import sample +from pathlib import Path + +from nonebot.permission import SUPERUSER + +from ATRI.service import Service +from ATRI.utils import request +from ATRI.exceptions import WriteFileError + + +CONSOLE_DIR = Path(".") / "data" / "database" / "console" +CONSOLE_DIR.mkdir(exist_ok=True) + + +class Console(Service): +    def __init__(self): +        Service.__init__( +            self, "控制台", "前端管理页面", True, main_cmd="/con", permission=SUPERUSER +        ) + +    @staticmethod +    async def get_host_ip(is_pub: bool): +        if is_pub: +            data = await request.get("https://ifconfig.me/ip") +            return data.text + +        s = None +        try: +            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) +            s.connect(('8.8.8.8', 80)) +            ip = s.getsockname()[0] +            return ip +        finally: +            if s: +                s.close() + +    @staticmethod +    def get_random_str(k: int) -> str: +        return "".join(sample(string.ascii_letters + string.digits, k)) + +    @staticmethod +    def get_auth_info() -> dict: +        df = CONSOLE_DIR / "data.json" +        if not df.is_file(): +            try: +                with open(df, "w", encoding="utf-8") as w: +                    w.write(json.dumps({})) +            except WriteFileError: +                raise WriteFileError("Writing file: " + str(df) + " failed!") + +        base_data: dict = json.loads(df.read_bytes()) +        data = base_data.get("data", None) +        if not data: +            return {"data": None} +        return data diff --git a/ATRI/plugins/console/driver/__init__.py b/ATRI/plugins/console/driver/__init__.py new file mode 100644 index 0000000..fe4398a --- /dev/null +++ b/ATRI/plugins/console/driver/__init__.py @@ -0,0 +1,56 @@ +from nonebot import get_driver +from nonebot.drivers.fastapi import Driver + +from fastapi.middleware.cors import CORSMiddleware + +from .view import ( +    handle_auther, +    handle_base_uri, +    handle_control_service, +    handle_edit_block, +    handle_get_block_list, +    handle_get_service_list, +    handle_runtime_info, +    handle_message_deal_info, +) + + +CONSOLE_API_URI = "/capi"  # base point +CONSOLE_API_AUTH_URI = "/capi/auth"  # 验证后台许可 +CONSOLE_API_RUNTIME_URI = "/capi/runtime"  # 获取运行占用信息 +CONSOLE_API_MESSAGE_URI = "/capi/message"  # 获取信息处理信息 + +CONSOLE_API_SERVICE_LIST_URI = "/capi/service/list"  # 获取服务列表 +CONSOLE_API_SERVICE_CONTROL_URI = "/capi/service/control"  # 对服务作出修改 + +CONSOLE_API_BLOCK_LIST_URI = "/capi/block/list"  # 获取封禁列表 +CONSOLE_API_BLOCK_EDIT_URI = "/capi/block/edit"  # 编辑封禁列表 + + +def register_routes(driver: Driver): +    app = driver.server_app + +    origins = ["*"] +    app.add_middleware( +        CORSMiddleware, +        allow_origins=origins, +        allow_credentials=True, +        allow_methods=["*"], +        allow_headers=["*"], +    ) + +    app.get(CONSOLE_API_URI)(handle_base_uri) +    app.get(CONSOLE_API_AUTH_URI)(handle_auther) +    app.get(CONSOLE_API_RUNTIME_URI)(handle_runtime_info) +    app.get(CONSOLE_API_MESSAGE_URI)(handle_message_deal_info) + +    app.get(CONSOLE_API_SERVICE_LIST_URI)(handle_get_service_list) +    app.get(CONSOLE_API_SERVICE_CONTROL_URI)(handle_control_service) + +    app.get(CONSOLE_API_BLOCK_LIST_URI)(handle_get_block_list) +    app.get(CONSOLE_API_BLOCK_EDIT_URI)(handle_edit_block) + + +def init(): +    driver = get_driver() +    register_routes(driver)  # type: ignore diff --git a/ATRI/plugins/console/driver/api.py b/ATRI/plugins/console/driver/api.py new file mode 100644 index 0000000..ba5405a --- /dev/null +++ b/ATRI/plugins/console/driver/api.py @@ -0,0 +1,193 @@ +import os +import json +import time +import psutil +from pathlib import Path +from datetime import datetime + +from ATRI.service import ServiceTools, SERVICES_DIR +from ATRI.exceptions import GetStatusError, ReadFileError, WriteFileError +from ..models import PlatformRuntimeInfo, BotRuntimeInfo, ServiceInfo + + +def get_processing_data() -> tuple: +    try: +        p_cpu = psutil.cpu_percent(interval=1) +        p_mem = psutil.virtual_memory().percent +        disk = psutil.disk_usage("/").percent +        inte_send = psutil.net_io_counters().bytes_sent / 1000000  # type: ignore +        inte_recv = psutil.net_io_counters().bytes_recv / 1000000  # type: ignore + +        process = psutil.Process(os.getpid()) +        b_cpu = process.cpu_percent(interval=1) +        b_mem = process.memory_percent(memtype="rss") + +        now = time.time() +        boot = psutil.boot_time() +        b = process.create_time() +        boot_time = str( +            datetime.utcfromtimestamp(now).replace(microsecond=0) +            - datetime.utcfromtimestamp(boot).replace(microsecond=0) +        ) +        bot_time = str( +            datetime.utcfromtimestamp(now).replace(microsecond=0) +            - datetime.utcfromtimestamp(b).replace(microsecond=0) +        ) +    except GetStatusError: +        raise GetStatusError("Getting runtime failed.") + +    if p_cpu > 90:  # type: ignore +        msg = "咱感觉有些头晕..." +        if p_mem > 90: +            msg = "咱感觉有点头晕并且有点累..." +    elif p_mem > 90: +        msg = "咱感觉有点累..." +    elif disk > 90: +        msg = "咱感觉身体要被塞满了..." +    else: +        msg = "アトリは、高性能ですから!" + +    return ( +        PlatformRuntimeInfo( +            stat_msg=msg, +            cpu_percent=str(p_cpu), +            mem_percent=p_mem, +            disk_percent=str(disk), +            inte_send=str(inte_send), +            inte_recv=str(inte_recv), +            boot_time=boot_time, +        ).dict(), +        BotRuntimeInfo( +            cpu_percent=str(b_cpu), mem_percent=str(b_mem), bot_run_time=bot_time +        ).dict(), +    ) + + +def get_service_list() -> dict: +    result = dict() + +    files = os.listdir(SERVICES_DIR) +    for f in files: +        # Thank you, MacOS +        if f == ".DS_Store": +            continue + +        serv_path = SERVICES_DIR / f +        data = json.loads(serv_path.read_bytes()) + +        serv_name = data["service"] +        serv_docs = data["docs"] +        serv_is_enabled = data["enabled"] +        serv_disable_user = data["disable_user"] +        serv_disable_group = data["disable_group"] + +        result[serv_name] = ServiceInfo( +            service_name=serv_name, +            service_docs=serv_docs, +            is_enabled=serv_is_enabled, +            disable_user=serv_disable_user, +            disable_group=serv_disable_group, +        ).dict() + +    return result + + +def control_service( +    serv_name: str, is_enab: int, enab_u: str, enab_g: str, disab_u: str, disab_g: str +) -> tuple: +    try: +        serv_data = ServiceTools().load_service(serv_name) +    except ReadFileError: +        return False, dict() + +    if is_enab != 1: +        if is_enab == 0: +            serv_data["enabled"] = False +        else: +            serv_data["enabled"] = True + +    if enab_u: +        if enab_u not in serv_data["disable_user"]: +            return False, {"msg": "Target not in list"} +        serv_data["disable_user"].remove(enab_u) + +    if enab_g: +        if enab_g not in serv_data["disable_user"]: +            return False, {"msg": "Target not in list"} +        serv_data["disable_group"].remove(enab_g) + +    if disab_u: +        if disab_u in serv_data["disable_user"]: +            return False, {"msg": "Target already exists in list"} +        serv_data["disable_user"].append(disab_u) + +    if disab_g: +        if disab_g in serv_data["disable_user"]: +            return False, {"msg": "Target already exists in list"} +        serv_data["disable_group"].append(disab_g) + +    try: +        ServiceTools().save_service(serv_data, serv_name) +    except WriteFileError: +        return False, dict() + +    return True, serv_data + + +MANEGE_DIR = Path(".") / "data" / "database" / "manege" + + +def get_block_list() -> dict: +    u_f = "block_user.json" +    path = MANEGE_DIR / u_f +    u_data = json.loads(path.read_bytes()) + +    g_f = "block_group.json" +    path = MANEGE_DIR / g_f +    g_data = json.loads(path.read_bytes()) + +    return {"user": u_data, "group": g_data} + + +def edit_block_list(is_enab: int, user_id: str, group_id: str) -> tuple: +    d = get_block_list() +    u_d = d["user"] +    g_d = d["group"] + +    now_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + +    if is_enab: +        if user_id: +            if user_id in u_d: +                return False, {"msg": "Target already exists in list"} +            u_d[user_id] = now_time + +        if group_id: +            if group_id in g_d: +                return False, {"msg": "Target already exists in list"} +            g_d[group_id] = now_time +    else: +        if user_id: +            if user_id not in u_d: +                return False, {"msg": "Target not in list"} +            del u_d[user_id] + +        if group_id: +            if group_id not in g_d: +                return False, {"msg": "Target not in list"} +            del g_d[group_id] + +    try: +        u_f = "block_user.json" +        path = MANEGE_DIR / u_f +        with open(path, "w", encoding="utf-8") as w: +            w.write(json.dumps(u_d)) + +        g_f = "block_group.json" +        path = MANEGE_DIR / g_f +        with open(path, "w", encoding="utf-8") as w: +            w.write(json.dumps(g_d)) +    except WriteFileError: +        return False, dict() + +    return True, {"user": u_d, "group": g_d} diff --git a/ATRI/plugins/console/driver/view.py b/ATRI/plugins/console/driver/view.py new file mode 100644 index 0000000..624e41f --- /dev/null +++ b/ATRI/plugins/console/driver/view.py @@ -0,0 +1,101 @@ +from ..data_source import Console +from ..listener import get_message_deal_info +from .api import ( +    control_service, +    edit_block_list, +    get_block_list, +    get_processing_data, +    get_service_list, +) + + +def auth_token(token: str) -> tuple: +    auth_data: dict = Console().get_auth_info() +    if not auth_data.get("token", None): +        return False, {"status": 500, "msg": "This bot is not create auth data yet."} +    _token = auth_data["token"] +    if token != _token: +        return False, {"status": 403, "msg": "Token error, please check again."} +    else: +        return True, {"status": 200, "msg": "OK"} + + +def handle_base_uri(): +    return {"status": 204, "msg": "This path just for console load."} + + +def handle_auther(token: str): +    auth, data = auth_token(token) +    return data if auth else data + + +def handle_runtime_info(token: str): +    auth, data = auth_token(token) +    if not auth: +        return data + +    plat, bot = get_processing_data() +    return {"status": 200, "data": {"platform": plat, "bot": bot}} + + +def handle_message_deal_info(token: str): +    auth, data = auth_token(token) +    if not auth: +        return data + +    return {"status": 200, "data": get_message_deal_info()} + + +def handle_get_service_list(token: str): +    auth, data = auth_token(token) +    if not auth: +        return data + +    return {"status": 200, "data": get_service_list()} + + +def handle_control_service( +    token: str, +    service: str, +    is_enabled: int = 1, +    enabled_user: str = str(), +    enabled_group: str = str(), +    disable_user: str = str(), +    disable_group: str = str(), +): +    auth, data = auth_token(token) +    if not auth: +        return data + +    is_ok, data = control_service( +        service, is_enabled, enabled_user, enabled_group, disable_user, disable_group +    ) +    if not is_ok: +        return {"status": 422, "msg": "Dealing service data failed"} + +    return {"status": 200, "data": data} + + +def handle_get_block_list(token: str): +    auth, data = auth_token(token) +    if not auth: +        return data + +    return {"status": 200, "data": get_block_list()} + + +def handle_edit_block( +    token: str, +    is_enabled: bool, +    user_id: str = str(), +    group_id: str = str(), +): +    auth, data = auth_token(token) +    if not auth: +        return data + +    is_ok, data = edit_block_list(is_enabled, user_id, group_id) +    if not is_ok: +        return {"status": 422, "msg": "Dealing block data failed"} + +    return {"status": 200, "data": data} diff --git a/ATRI/plugins/console/listener.py b/ATRI/plugins/console/listener.py new file mode 100644 index 0000000..0ae5c25 --- /dev/null +++ b/ATRI/plugins/console/listener.py @@ -0,0 +1,49 @@ +from typing import Optional + +from nonebot.message import run_postprocessor + +from ATRI.utils.apscheduler import scheduler +from .models import MessageDealerInfo + + +recv_msg = int() +deal_msg = int() +failed_deal_msg = int() + +total_r_m = int() +total_d_m = int() +total_f_m = int() + + +@run_postprocessor +async def _(exception: Optional[Exception]): +    global recv_msg, deal_msg, failed_deal_msg, total_r_m, total_d_m, total_f_m + +    if exception: +        failed_deal_msg += 1 +        total_f_m += 1 +    else: +        deal_msg += 1 +        total_d_m += 1 + +    recv_msg += 1 +    total_r_m += 1 + + +def get_message_deal_info() -> dict: +    return MessageDealerInfo( +        recv_msg=str(recv_msg), +        deal_msg=str(deal_msg), +        failed_deal_msg=str(failed_deal_msg), +        total_r_m=str(total_r_m), +        total_d_m=str(total_d_m), +        total_f_m=str(total_f_m), +    ).dict() + + +@scheduler.scheduled_job("interval", name="信息数据重置", seconds=15, misfire_grace_time=1)  # type: ignore +async def _(): +    global recv_msg, deal_msg, failed_deal_msg +    recv_msg = int() +    deal_msg = int() +    failed_deal_msg = int() diff --git a/ATRI/plugins/console/models.py b/ATRI/plugins/console/models.py new file mode 100644 index 0000000..973d05d --- /dev/null +++ b/ATRI/plugins/console/models.py @@ -0,0 +1,40 @@ +from pydantic import BaseModel + + +class AuthData(BaseModel): +    ip: str +    port: str +    token: str + + +class PlatformRuntimeInfo(BaseModel): +    stat_msg: str +    cpu_percent: str +    mem_percent: str +    disk_percent: str +    inte_send: str +    inte_recv: str +    boot_time: str + + +class BotRuntimeInfo(BaseModel): +    cpu_percent: str +    mem_percent: str +    bot_run_time: str + + +class MessageDealerInfo(BaseModel): +    recv_msg: str +    deal_msg: str +    failed_deal_msg: str +    total_r_m: str +    total_d_m: str +    total_f_m: str + + +class ServiceInfo(BaseModel): +    service_name: str +    service_docs: str +    is_enabled: bool +    disable_user: list +    disable_group: list diff --git a/ATRI/plugins/kimo/__init__.py b/ATRI/plugins/kimo/__init__.py index b1c5698..31ecab1 100644 --- a/ATRI/plugins/kimo/__init__.py +++ b/ATRI/plugins/kimo/__init__.py @@ -21,7 +21,8 @@ async def _chat(event: MessageEvent):      msg = str(event.message)      repo = await Kimo().deal(msg, user_id)      try: -        await kimo.finish(repo) +        if repo: +            await kimo.send(repo)      except Exception:          return diff --git a/ATRI/plugins/kimo/data_source.py b/ATRI/plugins/kimo/data_source.py index 6dd448a..fb3f3f3 100644 --- a/ATRI/plugins/kimo/data_source.py +++ b/ATRI/plugins/kimo/data_source.py @@ -8,7 +8,7 @@ from ATRI.service import Service  from ATRI.rule import to_bot, is_in_service  from ATRI.log import logger as log  from ATRI.utils import request -from ATRI.exceptions import ReadFileError, WriteError +from ATRI.exceptions import ReadFileError, WriteFileError  CHAT_PATH = Path(".") / "data" / "database" / "kimo" @@ -39,8 +39,8 @@ class Kimo(Service):                  with open(path, "w", encoding="utf-8") as w:                      w.write(json.dumps(data, indent=4))                  log.info("生成完成") -            except WriteError: -                raise WriteError("Writing kimo words failed!") +            except WriteFileError: +                raise WriteFileError("Writing kimo words failed!")      @classmethod      async def _load_data(cls) -> dict: diff --git a/ATRI/plugins/polaroid/data_source.py b/ATRI/plugins/polaroid/data_source.py index 82e0337..d4df4fe 100644 --- a/ATRI/plugins/polaroid/data_source.py +++ b/ATRI/plugins/polaroid/data_source.py @@ -4,7 +4,7 @@ from ATRI.service import Service  from ATRI.rule import is_in_service  from ATRI.utils import request  from ATRI.log import logger as log -from ATRI.exceptions import RequestError, WriteError +from ATRI.exceptions import RequestError, WriteFileError  from .image_dealer import image_dealer @@ -52,8 +52,8 @@ async def init_source():                  with open(path, "wb") as w:                      w.write(data.read())                  log.info("所需资源装载完成") -            except WriteError: -                raise WriteError("装载资源失败") +            except WriteFileError: +                raise WriteFileError("装载资源失败")  loop = asyncio.get_event_loop() diff --git a/ATRI/plugins/setu/tf_dealer.py b/ATRI/plugins/setu/tf_dealer.py index f966636..58b4337 100644 --- a/ATRI/plugins/setu/tf_dealer.py +++ b/ATRI/plugins/setu/tf_dealer.py @@ -13,7 +13,7 @@ import tensorflow as tf  from ATRI.log import logger as log  from ATRI.utils import request -from ATRI.exceptions import RequestError, WriteError +from ATRI.exceptions import RequestError, WriteFileError  SETU_PATH = Path(".") / "data" / "database" / "setu" @@ -60,8 +60,8 @@ async def detect_image(url: str, file_size: int) -> list:          path = TEMP_PATH / f"{file_name}.jpg"          with open(path, "wb") as f:              f.write(req.read()) -    except WriteError: -        raise WriteError("Writing file failed!") +    except WriteFileError: +        raise WriteFileError("Writing file failed!")      await init_module()      model_path = str((SETU_PATH / "nsfw.tflite").absolute()) @@ -107,8 +107,8 @@ async def init_module():              with open(path, "wb") as w:                  w.write(data.read())              log.info("模型装载完成") -        except WriteError: -            raise WriteError("NSFW TF module init failed!") +        except WriteFileError: +            raise WriteFileError("NSFW TF module init failed!")  loop = asyncio.get_event_loop() diff --git a/ATRI/plugins/status/__init__.py b/ATRI/plugins/status/__init__.py index 05ece79..ae7702a 100644 --- a/ATRI/plugins/status/__init__.py +++ b/ATRI/plugins/status/__init__.py @@ -1,21 +1,22 @@ +from ATRI.log import logger as log  from ATRI.utils.apscheduler import scheduler -from .data_source import IsSurvive +from .data_source import Status -ping = IsSurvive().on_command("/ping", "检测bot简单信息处理速度") +ping = Status().on_command("/ping", "检测bot简单信息处理速度")  @ping.handle() -async def _ping(): -    await ping.finish(IsSurvive.ping()) +async def _(): +    await ping.finish(Status.ping()) -status = IsSurvive().on_command("/status", "查看运行资源占用") +status = Status().on_command("/status", "查看运行资源占用")  @status.handle() -async def _status(): -    msg, _ = IsSurvive.get_status() +async def _(): +    msg, _ = Status.get_status()      await status.finish(msg) @@ -23,7 +24,8 @@ info_msg = "アトリは高性能ですから!"  @scheduler.scheduled_job("interval", name="状态检查", minutes=10, misfire_grace_time=15)  # type: ignore -async def _check_runtime(): -    msg, stat = IsSurvive().get_status() +async def _(): +    log.info("开始检查资源消耗...") +    msg, stat = Status().get_status()      if not stat:          await status.finish(msg) diff --git a/ATRI/plugins/status/data_source.py b/ATRI/plugins/status/data_source.py index 3595c0c..ddd3568 100644 --- a/ATRI/plugins/status/data_source.py +++ b/ATRI/plugins/status/data_source.py @@ -1,9 +1,9 @@ +import os  import time  import psutil  from datetime import datetime  from ATRI.service import Service -from ATRI.log import logger as log  from ATRI.rule import is_in_service  from ATRI.exceptions import GetStatusError @@ -11,19 +11,20 @@ from ATRI.exceptions import GetStatusError  _status_msg = """  > Status Overview -[CPU: {cpu}%] -[Memory: {mem}%] -[Disk usage: {disk}%] +[CPU: {p_cpu}% of {b_cpu}] +[Memory: {p_mem}% of {b_mem}] +[Disk usage: {p_disk}%]  [Net sent: {inteSENT}MB]  [Net recv: {inteRECV}MB] -[Runtime: {up_time}] +[Bot runtime: {bot_time}] +[Platform runtime: {boot_time}]  {msg}  """.strip() -class IsSurvive(Service): +class Status(Service):      def __init__(self):          Service.__init__(self, "状态", "检查自身状态", rule=is_in_service("状态")) @@ -32,21 +33,29 @@ class IsSurvive(Service):          return "I'm fine."      @staticmethod -    def get_status(): -        log.info("开始检查资源消耗...") +    def get_status() -> tuple:          try:              cpu = psutil.cpu_percent(interval=1)              mem = psutil.virtual_memory().percent              disk = psutil.disk_usage("/").percent -            inteSENT = psutil.net_io_counters().bytes_sent / 1000000  # type: ignore -            inteRECV = psutil.net_io_counters().bytes_recv / 1000000  # type: ignore +            inte_send = psutil.net_io_counters().bytes_sent / 1000000  # type: ignore +            inte_recv = psutil.net_io_counters().bytes_recv / 1000000  # type: ignore + +            process = psutil.Process(os.getpid()) +            b_cpu = process.cpu_percent(interval=1) +            b_mem = process.memory_percent(memtype="rss")              now = time.time()              boot = psutil.boot_time() -            up_time = str( +            b = process.create_time() +            boot_time = str(                  datetime.utcfromtimestamp(now).replace(microsecond=0)                  - datetime.utcfromtimestamp(boot).replace(microsecond=0)              ) +            bot_time = str( +                datetime.utcfromtimestamp(now).replace(microsecond=0) +                - datetime.utcfromtimestamp(b).replace(microsecond=0) +            )          except GetStatusError:              raise GetStatusError("Failed to get status.") @@ -64,16 +73,18 @@ class IsSurvive(Service):              msg = "咱感觉身体要被塞满了..."              is_ok = False          else: -            log.info("资源占用正常")              is_ok = True          msg0 = _status_msg.format( -            cpu=cpu, -            mem=mem, -            disk=disk, -            inteSENT=inteSENT, -            inteRECV=inteRECV, -            up_time=up_time, +            p_cpu=cpu, +            p_mem=mem, +            p_disk=disk, +            b_cpu=f"{b_cpu}%", +            b_mem="%.1f%%" % b_mem, +            inteSENT=inte_send, +            inteRECV=inte_recv, +            bot_time=bot_time, +            boot_time=boot_time,              msg=msg,          ) diff --git a/ATRI/service.py b/ATRI/service.py index c9509df..c52f5ca 100644 --- a/ATRI/service.py +++ b/ATRI/service.py @@ -17,7 +17,7 @@ from nonebot.typing import (  )  from nonebot.rule import Rule, command, keyword, regex -from ATRI.exceptions import ReadFileError, WriteError +from ATRI.exceptions import ReadFileError, WriteFileError  SERVICE_DIR = Path(".") / "data" / "service" @@ -101,8 +101,8 @@ class Service:          try:              with open(path, "w", encoding="utf-8") as w:                  w.write(json.dumps(data.dict(), indent=4)) -        except WriteError: -            raise WriteError("Write service info failed!") +        except WriteFileError: +            raise WriteFileError("Write service info failed!")      def save_service(self, service_data: dict, service: str) -> None:          if not service: @@ -304,9 +304,9 @@ class ServiceTools(object):          path = SERVICES_DIR / f"{service}.json"          if not path.is_file():              raise ReadFileError( -                f"Can't find service: ({service}) file.\n" -                "Please delete all file in data/service/services.\n" -                "Next reboot bot." +                f"Can't find service file: {service}\n" +                "Please delete all file in data/service/services\n" +                "And reboot bot."              )          with open(path, "w", encoding="utf-8") as w: @@ -317,9 +317,9 @@ class ServiceTools(object):          path = SERVICES_DIR / f"{service}.json"          if not path.is_file():              raise ReadFileError( -                f"Can't find service: ({service}) file.\n" -                "Please delete all file in data/service/services.\n" -                "Next reboot bot." +                f"Can't find service file: {service}\n" +                "Please delete all file in data/service/services\n" +                "And reboot bot."              )          with open(path, "r", encoding="utf-8") as r: | 
