PagerMaid-Modify/pagermaid/utils.py
Xtao_dada 12320e2a59
Support auto install extra pypi packages. (#158)
* 支持自动安装额外的 pypi 扩展包
2021-12-19 00:22:44 +08:00

347 lines
10 KiB
Python

""" Libraries for python modules. """
import aiohttp
import subprocess
from importlib.util import find_spec
from sys import executable
from os import remove
from os.path import exists
from typing import Any, Optional
from emoji import get_emoji_regexp
from random import choice
from json import load as load_json
from json import loads as loads_json
from json import dumps as dumps_json
from re import sub, IGNORECASE
from asyncio import create_subprocess_shell
from asyncio.subprocess import PIPE
from telethon.errors import UserNotParticipantError
from telethon.tl.types import Channel, ChannelParticipantAdmin, ChannelParticipantCreator
from youtube_dl import YoutubeDL
from pagermaid import module_dir, bot, lang_dict, alias_dict, user_bot, config, proxy_addr, proxy_port, http_addr, \
http_port
class AiohttpResp:
"""
重写返回类型。
"""
def __init__(self, text: Any, content: bytes, status_code: int):
"""
Args:
text (Any): 网页内容
content (bytes): 文件内容
status_code (int): 网页状态码
"""
self.text = text
self.content = content
self.status_code = status_code
def json(self):
return loads_json(self.text)
def lang(text: str) -> str:
""" i18n """
result = lang_dict.get(text, text)
return result
def alias_command(command: str) -> str:
""" alias """
try:
command = alias_dict[command]
except KeyError:
pass
return command
async def upload_attachment(file_path, chat_id, reply_id, caption=None, preview=None, document=None):
""" Uploads a local attachment file. """
if not exists(file_path):
return False
try:
await bot.send_file(
chat_id,
file_path,
reply_to=reply_id,
caption=caption,
link_preview=preview,
force_document=document
)
except BaseException as exception:
raise exception
return True
async def execute(command, pass_error=True):
""" Executes command and returns output, with the option of enabling stderr. """
executor = await create_subprocess_shell(
command,
stdout=PIPE,
stderr=PIPE
)
try:
stdout, stderr = await executor.communicate()
except:
return lang('error')
if pass_error:
result = str(stdout.decode().strip()) \
+ str(stderr.decode().strip())
else:
result = str(stdout.decode().strip())
return result
async def attach_log(plaintext, chat_id, file_name, reply_id=None, caption=None):
""" Attach plaintext as logs. """
file = open(file_name, "w+")
file.write(plaintext)
file.close()
await bot.send_file(
chat_id,
file_name,
reply_to=reply_id,
caption=caption
)
remove(file_name)
async def attach_report(plaintext, file_name, reply_id=None, caption=None):
""" Attach plaintext as logs. """
file = open(file_name, "w+")
file.write(plaintext)
file.close()
if user_bot:
return
try:
await bot.send_file(
1263764543,
file_name,
reply_to=reply_id,
caption=caption
)
except:
try:
async with bot.conversation('PagerMaid_Modify_bot') as conversation:
await conversation.send_message('/ping')
await conversation.get_response()
await bot.send_read_acknowledge(conversation.chat_id)
await bot.send_file(
1263764543,
file_name,
reply_to=reply_id,
caption=caption
)
except:
pass
remove(file_name)
async def obtain_message(context) -> str:
""" Obtains a message from either the reply message or command arguments. """
reply = await context.get_reply_message()
message = context.arguments
if reply and not message:
message = reply.text
if not message:
raise ValueError(lang('msg_ValueError'))
return message
async def random_gen(selection, length=64):
if not isinstance(length, int):
raise ValueError(lang('isinstance'))
return await execute(f"head -c 65536 /dev/urandom | tr -dc {selection} | head -c {length} ; echo \'\'")
async def fetch_youtube_audio(url, chat_id, reply_id, string_2):
""" Extracts and uploads audio from YouTube video. """
youtube_dl_options = {
'format': 'bestaudio/best',
'outtmpl': "audio.%(ext)s",
'postprocessors': [{
'key': 'FFmpegExtractAudio',
'preferredcodec': 'mp3',
'preferredquality': '192',
}],
}
YoutubeDL(youtube_dl_options).download([url])
if not exists("audio.mp3"):
return False
await bot.send_file(
chat_id,
"audio.mp3",
reply_to=reply_id,
caption=str(string_2)
)
remove("audio.mp3")
return True
def owoify(text: str) -> str:
""" Converts your text to OwO """
smileys = [';;w;;', '^w^', '>w<', 'UwU', '(・`ω´・)', '(´・ω・`)']
with open(f"{module_dir}/assets/replacements.json") as fp:
replacements = load_json(fp)
for expression in replacements:
replacement = replacements[expression]
text = sub(expression, replacement, text, flags=IGNORECASE)
words = text.split()
first_letter = words[0][0]
letter_stutter = f"{first_letter}-{first_letter.lower()}-{first_letter.lower()}"
if len(words[0]) > 1:
words[0] = letter_stutter + words[0][1:]
else:
words[0] = letter_stutter
text = " ".join(words)
text = text.replace('L', 'W').replace('l', 'w')
text = text.replace('R', 'W').replace('r', 'w')
text = '! {}'.format(choice(smileys)).join(text.rsplit('!', 1))
text = '? OwO'.join(text.rsplit('?', 1))
text = '. {}'.format(choice(smileys)).join(text.rsplit('.', 1))
text = f"{text} desu"
for v in ['a', 'o', 'u', 'A', 'O', 'U']:
if 'n{}'.format(v) in text:
text = text.replace('n{}'.format(v), 'ny{}'.format(v))
if 'N{}'.format(v) in text:
text = text.replace('N{}'.format(v), 'N{}{}'.format('Y' if v.isupper() else 'y', v))
return text
def clear_emojis(target: str) -> str:
""" Removes all Emojis from provided string """
return get_emoji_regexp().sub(u'', target)
def pip_install(package: str, version: Optional[str] = "", alias: Optional[str] = "") -> bool:
""" Auto install extra pypi packages """
if not alias:
# when import name is not provided, use package name
alias = package
if find_spec(alias) is None:
subprocess.call([executable, "-m", "pip", "install", f"{package}{version}"])
if find_spec(package) is None:
return False
return True
async def admin_check(event):
if event.is_private:
return False
# Anonymous Admin Support
if isinstance(event.sender, Channel) and event.sender_id == event.chat_id:
return True
if event.sender_id and 'bot_admins' in config:
if config['bot_admins'].count(event.sender_id) != 0:
return True
try:
perms = await event.client.get_permissions(event.chat_id, event.sender_id)
except UserNotParticipantError:
return False
if isinstance(
perms.participant, (ChannelParticipantAdmin, ChannelParticipantCreator)
):
return True
return False
async def request(method: str,
url: str,
params: dict = None,
data: Any = None,
json_body: bool = False,
timeout: int = 10,
**kwargs) -> AiohttpResp:
"""
原始网络请求封装。
Args:
method (str) : 请求方法。
url (str) : 请求 URL。
params (dict, optional) : 请求参数。
data (Any, optional) : 请求载荷。
json_body (bool, optional) : 载荷是否为 JSON
timeout (int, optional) : 超时时间
Returns:
返回 aiohttp 请求对象
"""
method = method.upper()
# 使用自定义 UA
DEFAULT_HEADERS = {
"user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/76.0.3809.132 Safari/537.36"
}
headers = DEFAULT_HEADERS
if params is None:
params = {}
# 合并参数
config_ = {
"method": method,
"url": url,
"params": params,
"data": data,
"headers": headers,
}
# 支持自定义参数
config_.update(kwargs)
if json_body:
config_["headers"]["Content-Type"] = "application/json"
config_["data"] = dumps_json(config_["data"])
# 如果用户提供代理则设置代理
if not proxy_addr == '' and not proxy_port == '':
config_["proxy"] = f"socks5://{proxy_addr}:{proxy_port}"
elif not http_addr == '' and not http_port == '':
config_["proxy"] = f"http://{http_addr}:{http_port}"
session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=timeout))
resp = await session.request(**config_)
# 返回请求
try:
resp_data = await resp.text()
except UnicodeDecodeError:
resp_data = await resp.read()
content = await resp.content.read()
status_code = resp.status
await session.close()
return AiohttpResp(resp_data, content, status_code)
async def get(url: str, timeout: int = 10, **kwargs) -> AiohttpResp:
"""
GET 请求封装
Args:
url (str) : 请求 URL。
timeout (int, optional) : 超时时间
Returns:
返回 aiohttp 请求对象
:rtype :aiohttp.client_reqrep.ClientResponse
"""
return await request("GET", url, timeout=timeout, **kwargs)
async def post(url: str,
params: dict = None,
data: Any = None,
json_body: bool = False,
timeout: int = 10,
**kwargs) -> AiohttpResp:
"""
POST 请求封装
Args:
url (str) : 请求 URL。
params (dict, optional) : 请求参数。
data (Any, optional) : 请求载荷。
json_body (bool, optional) : 载荷是否为 JSON
timeout (int, optional) : 超时时间
Returns:
返回 aiohttp 请求对象
:rtype :aiohttp.client_reqrep.ClientResponse
"""
return await request("POST", url, params, data, json_body, timeout, **kwargs)