PamGram/core/base/mtproto.py
2022-09-11 15:41:06 +08:00

63 lines
2.1 KiB
Python

import os
from typing import Optional
from urllib.parse import urlparse
import aiofiles
from utils.log import logger
try:
from pyrogram import Client
from pyrogram.session import session
session.log.debug = lambda *args, **kwargs: None # 关闭日记
PYROGRAM_AVAILABLE = True
except ImportError:
PYROGRAM_AVAILABLE = False
from core.bot import bot
from core.service import Service
class MTProto(Service):
async def get_session(self):
async with aiofiles.open(self.session_path, mode='r') as f:
return await f.read()
async def set_session(self, b: str):
async with aiofiles.open(self.session_path, mode='w+') as f:
await f.write(b)
def session_exists(self):
return os.path.exists(self.session_path)
def __init__(self):
self.name = "PaimonBot"
current_dir = os.getcwd()
self.session_path = os.path.join(current_dir, "paimon.session")
self.client: Optional[Client] = None
self.proxy: Optional[dict] = None
http_proxy = os.environ.get("HTTP_PROXY")
if http_proxy is not None:
http_proxy_url = urlparse(http_proxy)
self.proxy = {"scheme": "http", "hostname": http_proxy_url.hostname, "port": http_proxy_url.port}
async def start(self): # pylint: disable=W0221
if not PYROGRAM_AVAILABLE:
logger.info("MTProto 服务需要的 pyrogram 模块未导入 本次服务 client 为 None")
return
if bot.config.mtproto.api_id is None:
logger.info("MTProto 服务需要的 api_id 未配置 本次服务 client 为 None")
return
if bot.config.mtproto.api_hash is None:
logger.info("MTProto 服务需要的 api_hash 未配置 本次服务 client 为 None")
return
self.client = Client(api_id=bot.config.mtproto.api_id, api_hash=bot.config.mtproto.api_hash, name=self.name,
bot_token=bot.config.bot_token, proxy=self.proxy)
await self.client.start()
async def stop(self): # pylint: disable=W0221
if self.client is not None:
await self.client.stop(block=False)