PagerMaid-Modify/pagermaid/modules/status.py
Xtao_dada 4268f3bd48
Support official CLI to run speedtest. (#113)
 支持官方 CLI 程序运行 speedtest。
2021-07-27 12:14:49 +08:00

395 lines
14 KiB
Python

""" PagerMaid module that contains utilities related to system status. """
from json import loads
from PIL import Image
from requests import get
from os import remove, popen
from datetime import datetime
from speedtest import distance, Speedtest, ShareResultsConnectFailure, ShareResultsSubmitFailure, NoMatchedServers, \
SpeedtestBestServerFailure, SpeedtestHTTPError
from telethon import functions
from platform import python_version, uname
from wordcloud import WordCloud
from telethon import version as telethon_version
from telethon.tl.types import User, Chat, Channel
from sys import platform
from re import sub, findall
from pathlib import Path
from pagermaid import log, config, redis_status
from pagermaid.utils import execute, upload_attachment
from pagermaid.listener import listener
from pagermaid.utils import lang, alias_command
DCs = {
1: "149.154.175.50",
2: "149.154.167.51",
3: "149.154.175.100",
4: "149.154.167.91",
5: "91.108.56.130"
}
@listener(is_plugin=False, outgoing=True, command=alias_command("sysinfo"),
description=lang('sysinfo_des'))
async def sysinfo(context):
""" Retrieve system information via neofetch. """
await context.edit(lang('sysinfo_loading'))
result = await execute("neofetch --config none --stdout")
await context.edit(f"`{result}`")
@listener(is_plugin=False, outgoing=True, command=alias_command("fortune"),
description=lang('fortune_des'))
async def fortune(context):
""" Reads a fortune cookie. """
result = await execute("fortune")
if result == "/bin/sh: fortune: command not found":
await context.edit(lang('fortune_not_exist'))
return
await context.edit(result)
@listener(is_plugin=False, outgoing=True, command=alias_command("fbcon"),
description=lang('fbcon_des'))
async def tty(context):
""" Screenshots a TTY and prints it. """
await context.edit(lang('fbcon_processing'))
reply_id = context.message.reply_to_msg_id
result = await execute("fbdump | convert - image.png")
if result == "/bin/sh: fbdump: command not found":
await context.edit(lang('fbcon_no_fbdump'))
remove("image.png")
return
if result == "/bin/sh: convert: command not found":
await context.edit(lang('fbcon_no_ImageMagick'))
remove("image.png")
return
if result == "Failed to open /dev/fb0: Permission denied":
await context.edit(lang('fbcon_no_permission'))
return
if not await upload_attachment("./image.png", context.chat_id, reply_id,
caption=lang('fbcon_caption'),
preview=False, document=False):
await context.edit(lang('fbcon_error'))
return
await context.delete()
try:
remove("./image.png")
except:
pass
await log("Screenshot of binded framebuffer console taken.")
@listener(is_plugin=False, outgoing=True, command=alias_command("status"),
description=lang('status_des'))
async def status(context):
database = lang('status_online') if redis_status() else lang('status_offline')
text = (f"**{lang('status_hint')}** \n"
f"{lang('status_name')}: `{uname().node}` \n"
f"{lang('status_platform')}: `{platform}` \n"
f"{lang('status_release')}: `{uname().release}` \n"
f"{lang('status_python')}: `{python_version()}` \n"
f"{lang('status_telethon')}: `{telethon_version.__version__}` \n"
f"{lang('status_db')}: `{database}`"
)
await context.edit(text)
@listener(is_plugin=False, outgoing=True, command=alias_command("stats"),
description=lang('stats_des'))
async def stats(context):
await context.edit(lang('stats_loading'))
u, g, s, c, b = 0, 0, 0, 0, 0
dialogs = await context.client.get_dialogs(
limit=None,
ignore_migrated=True
)
for d in dialogs:
current_entity = d.entity
if type(current_entity) is User:
if current_entity.bot:
b += 1
else:
u += 1
elif type(current_entity) is Chat:
g += 1
elif type(current_entity) is Channel:
if current_entity.broadcast:
c += 1
else:
s += 1
text = (f"**{lang('stats_hint')}** \n"
f"{lang('stats_dialogs')}: `{len(dialogs)}` \n"
f"{lang('stats_private')}: `{u}` \n"
f"{lang('stats_group')}: `{g}` \n"
f"{lang('stats_supergroup')}: `{s}` \n"
f"{lang('stats_channel')}: `{c}` \n"
f"{lang('stats_bot')}: `{b}`"
)
await context.edit(text)
@listener(is_plugin=False, outgoing=True, command=alias_command("speedtest"),
description=lang('speedtest_des'))
async def speedtest(context):
""" Tests internet speed using speedtest. """
try:
speed_test_path = config['speed_test_path']
except KeyError:
speed_test_path = ''
if not speed_test_path == '':
server = None
if len(context.parameter) == 1:
try:
server = int(context.parameter[0])
except ValueError:
await context.edit(lang('arg_error'))
return
speed_test_path += ' -f json'
if server:
speed_test_path += f' -s {server}'
await context.edit(lang('speedtest_processing'))
result = await execute(f'{speed_test_path}')
result = loads(result)
if result['type'] == 'log':
await context.edit(f"{result['level'].upper()}:{result['message']}")
elif result['type'] == 'result':
des = (
f"**Speedtest** \n"
f"Server: `{result['server']['name']} - "
f"{result['server']['location']}` \n"
f"Host: `{result['server']['host']}` \n"
f"Upload: `{unit_convert(result['upload']['bandwidth'] * 8)}` \n"
f"Download: `{unit_convert(result['download']['bandwidth'] * 8)}` \n"
f"Latency: `{result['ping']['latency']}` \n"
f"Jitter: `{result['ping']['jitter']}` \n"
f"Timestamp: `{result['timestamp']}`"
)
# 开始处理图片
data = get(f"{result['result']['url']}.png").content
with open('speedtest.png', mode='wb') as f:
f.write(data)
try:
img = Image.open('speedtest.png')
c = img.crop((17, 11, 727, 389))
c.save('speedtest.png')
except:
pass
try:
await context.client.send_file(context.chat_id, 'speedtest.png', caption=des)
except:
pass
try:
remove('speedtest.png')
except:
pass
await context.delete()
else:
await context.edit(result)
return
try:
test = Speedtest()
except SpeedtestHTTPError:
await context.edit(lang('speedtest_ConnectFailure'))
return
server, server_json = [], False
if len(context.parameter) == 1:
try:
server = [int(context.parameter[0])]
except ValueError:
await context.edit(lang('arg_error'))
return
elif len(context.parameter) >= 2:
try:
temp_json = findall(r'{(.*?)}', context.text.replace("'", '"'))
if len(temp_json) == 1:
server_json = loads("{" + temp_json[0] + "}")
server_json['d'] = distance(test.lat_lon, (float(server_json['lat']), float(server_json['lon'])))
test.servers = [server_json]
else:
await context.edit(lang('arg_error'))
return
except:
pass
await context.edit(lang('speedtest_processing'))
try:
if len(server) == 0:
if not server_json:
test.get_best_server()
else:
test.get_best_server(servers=test.servers)
else:
test.get_servers(servers=server)
except (SpeedtestBestServerFailure, NoMatchedServers) as e:
await context.edit(lang('speedtest_ServerFailure'))
return
test.download()
test.upload()
try:
test.results.share()
except (ShareResultsConnectFailure, ShareResultsSubmitFailure) as e:
await context.edit(lang('speedtest_ConnectFailure'))
return
result = test.results.dict()
des = (
f"**Speedtest** \n"
f"Server: `{result['server']['name']} - "
f"{result['server']['cc']}` \n"
f"Sponsor: `{result['server']['sponsor']}` \n"
f"Upload: `{unit_convert(result['upload'])}` \n"
f"Download: `{unit_convert(result['download'])}` \n"
f"Latency: `{result['ping']}` \n"
f"Timestamp: `{result['timestamp']}`"
)
# 开始处理图片
data = get(result['share']).content
with open('speedtest.png', mode='wb') as f:
f.write(data)
try:
img = Image.open('speedtest.png')
c = img.crop((17, 11, 727, 389))
c.save('speedtest.png')
except:
pass
try:
await context.client.send_file(context.chat_id, 'speedtest.png', caption=des)
except:
return
try:
remove('speedtest.png')
except:
pass
await context.delete()
@listener(is_plugin=False, outgoing=True, command=alias_command("connection"),
description=lang('connection_des'))
async def connection(context):
""" Displays connection information between PagerMaid and Telegram. """
datacenter = await context.client(functions.help.GetNearestDcRequest())
await context.edit(
f"**{lang('connection_hint')}** \n"
f"{lang('connection_country')}: `{datacenter.country}` \n"
f"{lang('connection_dc')}: `{datacenter.this_dc}` \n"
f"{lang('connection_nearest_dc')}: `{datacenter.nearest_dc}`"
)
@listener(is_plugin=False, outgoing=True, command=alias_command("pingdc"),
description=lang('pingdc_des'))
async def pingdc(context):
""" Ping your or other data center's IP addresses. """
data = []
for dc in range(1, 6):
result = await execute(f"ping -c 1 {DCs[dc]} | awk -F '/' " + "'END {print $5}'")
data.append(result)
await context.edit(
f"{lang('pingdc_1')}: `{data[0]}ms`\n"
f"{lang('pingdc_2')}: `{data[1]}ms`\n"
f"{lang('pingdc_3')}: `{data[2]}ms`\n"
f"{lang('pingdc_4')}: `{data[3]}ms`\n"
f"{lang('pingdc_5')}: `{data[4]}ms`"
)
@listener(is_plugin=False, outgoing=True, command=alias_command("ping"),
description=lang('ping_des'))
async def ping(context):
""" Calculates latency between PagerMaid and Telegram. """
start = datetime.now()
await context.edit("Pong!")
end = datetime.now()
duration = (end - start).microseconds / 1000
await context.edit(f"Pong!|{duration}")
@listener(is_plugin=False, outgoing=True, command=alias_command("topcloud"),
description=lang('topcloud_des'))
async def topcloud(context):
""" Generates a word cloud of resource-hungry processes. """
await context.edit(lang('topcloud_processing'))
command_list = []
if not Path('/usr/bin/top').is_symlink():
output = str(await execute("top -b -n 1")).split("\n")[7:]
else:
output = str(await execute("top -b -n 1")).split("\n")[4:]
for line in output[:-1]:
line = sub(r'\s+', ' ', line).strip()
fields = line.split(" ")
try:
if fields[11].count("/") > 0:
command = fields[11].split("/")[0]
else:
command = fields[11]
cpu = float(fields[8].replace(",", "."))
mem = float(fields[9].replace(",", "."))
if command != "top":
command_list.append((command, cpu, mem))
except BaseException:
pass
command_dict = {}
for command, cpu, mem in command_list:
if command in command_dict:
command_dict[command][0] += cpu
command_dict[command][1] += mem
else:
command_dict[command] = [cpu + 1, mem + 1]
resource_dict = {}
for command, [cpu, mem] in command_dict.items():
resource_dict[command] = (cpu ** 2 + mem ** 2) ** 0.5
width, height = None, None
try:
width, height = ((popen("xrandr | grep '*'").read()).split()[0]).split("x")
width = int(width)
height = int(height)
except BaseException:
pass
if not width or not height:
width = int(config['width'])
height = int(config['height'])
background = config['background']
margin = int(config['margin'])
try:
cloud = WordCloud(
background_color=background,
width=width - 2 * int(margin),
height=height - 2 * int(margin)
).generate_from_frequencies(resource_dict)
except ValueError:
await context.edit(lang('run_error'))
return
cloud.to_file("cloud.png")
await context.edit(lang('highlight_uploading'))
await context.client.send_file(
context.chat_id,
"cloud.png",
reply_to=None,
caption=lang('topcloud_caption')
)
remove("cloud.png")
await context.delete()
await log(lang('topcloud_success'))
def unit_convert(byte):
""" Converts byte into readable formats. """
power = 1000
zero = 0
units = {
0: '',
1: 'Kb/s',
2: 'Mb/s',
3: 'Gb/s',
4: 'Tb/s'}
while byte > power:
byte /= power
zero += 1
return f"{round(byte, 2)} {units[zero]}"