209 lines
8.2 KiB
Python
209 lines
8.2 KiB
Python
import datetime
|
||
import time
|
||
import connector
|
||
import telegram
|
||
from telegram.ext import CommandHandler, MessageHandler, Filters
|
||
from config import TOKEN, LIMIT_COUNT, EXCLUSIVE_MODE, RANK_COMMAND_MODE, OWNER, EXCLUSIVE_LIST, CHANNEL, HELP
|
||
import schedule
|
||
from task import add_task, add_user_task
|
||
|
||
bot = telegram.Bot(token=TOKEN)
|
||
|
||
|
||
def start(update, context):
|
||
# 限制不为群组
|
||
chat_type = update.effective_chat.type
|
||
if chat_type == "supergroup":
|
||
return
|
||
try:
|
||
connector.get_connection().keys()
|
||
print('进入 start 函数')
|
||
update.message.reply_text(
|
||
HELP,
|
||
parse_mode='HTML'
|
||
)
|
||
except Exception as e:
|
||
print(e)
|
||
print('进入start函数')
|
||
if update.effective_user.id == OWNER:
|
||
update.message.reply_text(f"系统故障，Redis连接失败，错误信息：\n{e}")
|
||
|
||
|
||
def ping(update, context):
|
||
# 限制不为群组
|
||
chat_type = update.effective_chat.type
|
||
if chat_type == "supergroup":
|
||
return
|
||
try:
|
||
connector.get_connection().keys()
|
||
print('进入 ping 函数')
|
||
update.message.reply_text(
|
||
'pong~',
|
||
)
|
||
except Exception as e:
|
||
print(e)
|
||
|
||
|
||
def rank(update, context):
|
||
try:
|
||
r = connector.get_connection()
|
||
chat_type = update.effective_chat.type
|
||
user_id = update.effective_user.id
|
||
chat_id = update.effective_message.chat_id
|
||
try:
|
||
username = update.effective_user.username
|
||
except Exception as e:
|
||
username = update.effective_user.id
|
||
# 限制为群组
|
||
if chat_type != "supergroup":
|
||
update.message.reply_text("此命令只有在群组中有效！")
|
||
return
|
||
if RANK_COMMAND_MODE:
|
||
try:
|
||
chat_member = bot.get_chat_member(chat_id, user_id)
|
||
status = chat_member["status"]
|
||
if user_id == 1087968824:
|
||
status = "administrator"
|
||
print("此用户在群组中身份为： {}".format(status))
|
||
if status == "creator" or status == "administrator":
|
||
print("用户权限正确")
|
||
else:
|
||
return
|
||
except Exception as e:
|
||
print(e)
|
||
print("获取用户身份失败")
|
||
if r.exists("{}_frequency_limit".format(chat_id)):
|
||
r.setrange("{}_frequency_limit".format(chat_id), 0, int(r.get("{}_frequency_limit".format(chat_id))) + 1)
|
||
else:
|
||
struct_time = time.localtime(time.time())
|
||
# 数据过期时间为当前小时的 59 分
|
||
ex_time = datetime.datetime(
|
||
struct_time.tm_year,
|
||
struct_time.tm_mon,
|
||
struct_time.tm_mday,
|
||
struct_time.tm_hour,
|
||
59
|
||
)
|
||
r.set("{}_frequency_limit".format(chat_id), 1)
|
||
r.expireat("{}_frequency_limit".format(chat_id), ex_time)
|
||
count = int(r.get("{}_frequency_limit".format(chat_id)))
|
||
if count > LIMIT_COUNT:
|
||
update.message.reply_text("该群组在这个小时内的生成配额已经用完，请稍后再试~")
|
||
return
|
||
add_task(chat_id)
|
||
print("群组: {}，用户: {}|{} 发起了主动触发请求".format(chat_id, username, user_id, ))
|
||
if not CHANNEL == 0:
|
||
ctext = f'#WORDCLOUD #APPLY \n' \
|
||
f'群组 ID：`{chat_id}`\n' \
|
||
f'用户 ID：`{user_id}`\n' \
|
||
f'执行操作：`主动生成词云`'
|
||
bot.send_message(chat_id=CHANNEL, text=ctext, parse_mode="Markdown")
|
||
update.message.reply_text("统计数据将在分析完毕后发送到当前群组，请稍等~")
|
||
except Exception as e:
|
||
print("主动触发任务失败，请检查")
|
||
print(e)
|
||
|
||
|
||
def stat(update, context):
|
||
try:
|
||
r = connector.get_connection()
|
||
chat_type = update.effective_chat.type
|
||
user_id = update.effective_user.id
|
||
chat_id = update.effective_message.chat_id
|
||
try:
|
||
username = update.effective_user.username
|
||
except Exception as e:
|
||
username = update.effective_user.id
|
||
# 限制为群组
|
||
if chat_type != "supergroup":
|
||
update.message.reply_text("此命令只有在群组中有效！")
|
||
return
|
||
if r.exists("{}_{}_frequency_limit".format(chat_id, user_id)):
|
||
r.setrange("{}_{}_frequency_limit".format(chat_id, user_id), 0,
|
||
int(r.get("{}_{}_frequency_limit".format(chat_id, user_id))) + 1)
|
||
else:
|
||
struct_time = time.localtime(time.time())
|
||
# 数据过期时间为当前小时的 59 分
|
||
ex_time = datetime.datetime(
|
||
struct_time.tm_year,
|
||
struct_time.tm_mon,
|
||
struct_time.tm_mday,
|
||
struct_time.tm_hour,
|
||
59
|
||
)
|
||
r.set("{}_{}_frequency_limit".format(chat_id, user_id), 1)
|
||
r.expireat("{}_{}_frequency_limit".format(chat_id, user_id), ex_time)
|
||
count = int(r.get("{}_{}_frequency_limit".format(chat_id, user_id)))
|
||
if count > LIMIT_COUNT:
|
||
update.message.reply_text(f"[您](tg://user?id={user_id})在这个小时内的生成配额已经用完，请稍后再试~")
|
||
return
|
||
add_user_task(chat_id, user_id)
|
||
print("群组: {}，用户: {} | {} 发起了主动触发请求".format(chat_id, username, user_id, ))
|
||
if not CHANNEL == 0:
|
||
ctext = f'#WORDCLOUD #APPLY #id{user_id} \n' \
|
||
f'群组 ID：`{chat_id}`\n' \
|
||
f'用户 ID：`{user_id}`\n' \
|
||
f'执行操作：`主动生成用户词云`'
|
||
bot.send_message(chat_id=CHANNEL, text=ctext, parse_mode="Markdown")
|
||
update.message.reply_text("统计数据将在分析完毕后发送到当前群组，请稍等~")
|
||
except Exception as e:
|
||
print("主动触发任务失败，请检查")
|
||
print(e)
|
||
|
||
|
||
def chat_content_exec(update, context):
|
||
try:
|
||
r = connector.get_connection()
|
||
text = update.message.text
|
||
chat_type = update.effective_chat.type
|
||
user_id = update.effective_user.id
|
||
chat_id = update.effective_message.chat_id
|
||
# 限制为群组
|
||
if chat_type != "supergroup":
|
||
return
|
||
# 限制文字长度不能超过80字
|
||
if len(text) > 80:
|
||
return
|
||
# 独享模式（仅授权群组可用）
|
||
if EXCLUSIVE_MODE and chat_id not in EXCLUSIVE_LIST:
|
||
print(chat_id + " 为未认证群组，取消入库")
|
||
return
|
||
user = update.message.from_user
|
||
firstname = str(user["first_name"])
|
||
name = firstname
|
||
print("\n---------------------------")
|
||
print("内容: " + text[:10])
|
||
print("群组类型: " + str(chat_type))
|
||
print("用户ID: " + str(user_id))
|
||
print("chat_id: " + str(chat_id))
|
||
if text.startswith('/') or '//' in text or text.startswith('-'):
|
||
print("这是一条指令或者链接信息，跳过")
|
||
return
|
||
else:
|
||
if text[-1] not in ["，", "。", "！", "：", "？", "!", "?", ",", ":", "."]:
|
||
r.append("{}_chat_content".format(chat_id), text + "。")
|
||
r.append("{}_{}_user_content".format(chat_id, user_id), text + "。")
|
||
else:
|
||
r.append("{}_chat_content".format(chat_id), text)
|
||
r.append("{}_{}_user_content".format(chat_id, user_id), text)
|
||
r.incrby("{}_total_message_amount".format(chat_id))
|
||
r.incrby("{}_{}_user_message_amount".format(chat_id, user_id))
|
||
r.hincrby("{}_user_message_amount".format(chat_id), name)
|
||
print("---------------------------")
|
||
except Exception as e:
|
||
print(e)
|
||
print("用户数据提取、入库错误")
|
||
|
||
|
||
def check_schedule():
|
||
while True:
|
||
schedule.run_pending()
|
||
time.sleep(1)
|
||
|
||
|
||
start_handler = CommandHandler('start', start)
|
||
ping_handler = CommandHandler('ping', ping)
|
||
rank_handler = CommandHandler('rank', rank)
|
||
stat_handler = CommandHandler('stat', stat)
|
||
chat_content_handler = MessageHandler(Filters.text, chat_content_exec)
|