refactor code

This commit is contained in:
hibiki
2021-07-30 21:21:51 +08:00
parent 2ad891aa1e
commit cc24822dca
165 changed files with 7815 additions and 8174 deletions
+260 -290
View File
@@ -1,252 +1,209 @@
import random
from nonebot import on_command, on_regex
from nonebot.permission import SUPERUSER
from services.log import logger
from models.sigin_group_user import SignGroupUser
from utils.utils import FreqLimiter, UserExistLimiter, is_number, get_message_text, get_message_imgs
from nonebot.exception import IgnoredException
from nonebot.message import run_preprocessor, run_postprocessor
from nonebot.matcher import Matcher
from typing import Optional, Type
from gino.exceptions import UninitializedError
from utils.utils import (
FreqLimiter,
UserExistLimiter,
is_number,
get_message_text,
get_message_imgs,
)
from nonebot.typing import T_State
from nonebot.adapters.cqhttp import Bot, MessageEvent, GroupMessageEvent, PrivateMessageEvent
from .data_source import get_setu, get_luoxiang, search_online_setu, get_setu_urls, \
check_r18_and_keyword, find_img_index, delete_img, add_img
from nonebot.adapters.cqhttp import (
Bot,
MessageEvent,
GroupMessageEvent,
PrivateMessageEvent,
Message,
Event,
)
from .data_source import (
get_setu_list,
get_luoxiang,
search_online_setu,
get_setu_urls,
find_img_index,
gen_message,
check_local_exists_or_download,
add_data_to_database,
)
from nonebot.adapters.cqhttp.exception import ActionFailed
from configs.config import ONLY_USE_LOCAL_SETU, WITHDRAW_SETU_TIME
from utils.message_builder import at
import re
from models.count_user import UserCount
from aiohttp.client_exceptions import ClientConnectorError
from configs.config import MAX_SETU_R_COUNT
from configs.path_config import TXT_PATH
import asyncio
try:
import ujson as json
except ModuleNotFoundError:
import json
__plugin_name__ = '色图'
__plugin_usage__ = f'''示例:
__plugin_name__ = "色图"
__plugin_usage__ = f"""示例:
1. 色图 (随机本地色图)
2. 色图r (随机在线十张r18涩图)
3. 色图 666 (本地色图id)
4. 色图 xx (在线搜索xx色图)
4. 色图 xx xx (在线搜索xx xx色图)
5. 色图r xx (搜索十张xx的r18涩图,注意空格)(仅私聊,每日限制5次)
6. 来n张涩图 (本地涩图连发)(1<=n<=9)
7. 来n张xx的涩图 (在线搜索xx涩图)(较慢,看网速)
注:【色图r每日提供{MAX_SETU_R_COUNT}次
注:
xx 为 tag,多余20取前20 示例:原神 黑丝
本地涩图没有r18!
联网搜索会较慢!
如果图片数量与数字不符,
原因1:网络不好,网线被拔QAQ
原因2:搜索到的总数小于数字
原因3:图太色或者小错误了】'''
原因3:图太色或者小错误了】"""
url = "https://api.lolicon.app/setu/"
_flmt = FreqLimiter(5)
_ulmt = UserExistLimiter()
path = "_setu/"
setu_data_list = []
withdraw_message_id = []
setu = on_command("色图", aliases={"涩图", "不够色", "来一发", "再来点"}, priority=5, block=True)
setu_reg = on_regex('(.*)[份|发|张|个|次|点](.*)[瑟|色|涩]图', priority=5, block=True)
find_setu = on_command("查色图", priority=5, block=True)
delete_setu = on_command('删除色图', aliases={'删除涩图'}, priority=5, block=True, permission=SUPERUSER)
upload_setu = on_command('上传色图', aliases={'上传涩图'}, priority=5, block=True, permission=SUPERUSER)
@setu.handle()
async def _(bot: Bot, event: GroupMessageEvent, state: T_State):
img_id = get_message_text(event.json())
if img_id in ['帮助']:
await setu.finish(__plugin_usage__)
impression = (await SignGroupUser.ensure(event.user_id, event.group_id)).impression
luox = get_luoxiang(impression)
if luox:
await setu.finish(luox)
if not _flmt.check(event.user_id):
await setu.finish('您冲得太快了,请稍候再冲', at_sender=True)
_flmt.start_cd(event.user_id)
if _ulmt.check(event.user_id):
await setu.finish(f"您有色图正在处理,请稍等")
_ulmt.set_True(event.user_id)
try:
setu_data = json.load(open(f'{TXT_PATH}/setu_data.json', 'r', encoding='utf8'))
except Exception:
setu_data = {}
setu_img, index = get_setu(img_id, setu_data)
if setu_img:
try:
await setu.send(setu_img)
except:
_ulmt.set_False(event.user_id)
await setu.finish('这张图色过头了,我自己看看就行了!', at_sender=True)
logger.info(
f"USER {event.user_id} GROUP {event.group_id} 发送色图 {index}.jpg 成功")
else:
msg = img_id
if msg.find('真寻') != -1:
_ulmt.set_False(event.user_id)
await setu.finish('咳咳咳,虽然我很可爱,但是我木有自己的色图~~~有的话记得发我一份呀')
keyword, r18, num = await check_r18_and_keyword(msg, event.user_id)
if r18 == 1:
_ulmt.set_False(event.user_id)
await setu.finish(random.choice([
"这种不好意思的东西怎么可能给这么多人看啦",
"羞羞脸!给我滚出克私聊!",
"变态变态变态变态大变态!"
]))
try:
urls, text_list, code = await get_setu_urls(keyword, num, r18=r18)
except ClientConnectorError:
_ulmt.set_False(event.user_id)
await setu.finish('网络失败了..别担心!正在靠运气上网!', at_sender=True)
else:
if code == 200:
for i in range(num):
try:
setu_img, index = await search_online_setu(urls[i])
if index == -1:
setu_img, index = get_setu('', setu_data, keyword)
if setu_img:
await setu.send(setu_img)
logger.error(f'色图下载失败 本地发送 {index}.jpg')
else:
await setu.send(text_list[i] + setu_img)
logger.info(
f"USER {event.user_id} GROUP {event.group_id}"
f" 发送在线色图 {keyword}.jpg 成功")
except Exception as e:
setu_img, index = get_setu('', setu_data, keyword)
if setu_img:
await setu.send(setu_img)
logger.error(f'色图发送错误 本地发送 {index}.jpg e:{e}')
else:
await setu.send(urls)
_ulmt.set_False(event.user_id)
@setu.handle()
async def _(bot: Bot, event: PrivateMessageEvent, state: T_State):
img_id = get_message_text(event.json())
if img_id in ['帮助']:
await setu.finish(__plugin_usage__)
if not _flmt.check(event.user_id):
await setu.finish('您冲得太快了,请稍候再冲', at_sender=True)
_flmt.start_cd(event.user_id)
if _ulmt.check(event.user_id):
await setu.finish(f"您有色图正在处理,请稍等")
_ulmt.set_True(event.user_id)
try:
setu_data = json.load(open(f'{TXT_PATH}/setu_data.json', 'r', encoding='utf8'))
except Exception:
setu_data = {}
setu_img, index = get_setu(img_id, setu_data)
if setu_img:
await setu.send(setu_img)
logger.info(
f"USER {event.user_id} GROUP private 发送色图 {index}.jpg 成功")
else:
msg = img_id
if msg.find('真寻') != -1:
_ulmt.set_False(event.user_id)
await setu.finish('咳咳咳,虽然我很可爱,但是我木有自己的色图~~~有的话记得发我一份呀')
keyword, r18, num = await check_r18_and_keyword(msg, event.user_id)
if r18 == 1:
if await UserCount.check_count(event.user_id, 'setu_r18', MAX_SETU_R_COUNT):
_ulmt.set_False(event.user_id)
await setu.finish('要节制啊,请明天再来...\n【每日提供5次】', at_sender=True)
else:
await UserCount.add_count(event.user_id, 'setu_r18', count=1)
try:
setu_data = json.load(open(f'{TXT_PATH}/r18_setu_data.json', 'r', encoding='utf8'))
except Exception:
setu_data = {}
try:
urls, text_list, code = await get_setu_urls(keyword, num, r18=r18)
except ClientConnectorError:
await UserCount.add_count(event.user_id, 'setu_r18', count=-1)
_ulmt.set_False(event.user_id)
await setu.finish('网络失败了..别担心!这次搜索不算数喔', at_sender=True)
else:
count = 0
if code == 200:
for i in range(num):
try:
setu_img, index = await search_online_setu(urls[i])
if index == -1:
if img_id == 'r':
setu_img, index = get_setu('', setu_data, keyword, True)
else:
setu_img, index = get_setu('', setu_data, keyword)
if setu_img:
await setu.send(setu_img)
logger.error(f'色图下载失败 本地发送 {index}.jpg')
else:
await setu.send(text_list[i] + setu_img)
logger.info(
f"USER {event.user_id} GROUP private"
f" 发送{'r18' if img_id == 'r' else ''}色图 {index}.jpg 成功")
except Exception as e:
if img_id == 'r':
img_setu, index = get_setu('', setu_data, keyword, True)
else:
img_setu, index = get_setu('', setu_data, keyword)
if img_setu:
await setu.send(img_setu)
logger.error(f'色图发送错误 本地发送 {index}.jpg e:{e}')
count += 1
if count > 6:
await setu.send('检测到下载惜败的图片过多,这次就不算数了,果咩..', at_sender=True)
await UserCount.add_count(event.user_id, 'setu_r18', count=-1)
else:
if code == 401:
if r18 == 1:
await UserCount.add_count(event.user_id, 'setu_r18', count=-1)
await setu.send(urls + ' 色图r次数返还!')
else:
await setu.send(urls, at_sender=True)
@run_preprocessor
async def _(matcher: Matcher, bot: Bot, event: MessageEvent, state: T_State):
if isinstance(event, MessageEvent):
if matcher.module == "send_setu":
if _ulmt.check(event.user_id):
if isinstance(event, GroupMessageEvent):
await bot.send_group_msg(
group_id=event.group_id,
message=Message(f"{at(event.user_id)}您有色图正在处理,请稍等....."),
)
else:
if r18 == 1:
await setu.send('这次不是小真寻的戳!色图r次数返还!', at_sender=True)
await UserCount.add_count(event.user_id, 'setu_r18', count=-1)
else:
await setu.send(urls, at_sender=True)
_ulmt.set_False(event.user_id)
await bot.send_private_msg(
user_id=event.user_id, message=f"您有色图正在处理,请稍等....."
)
raise IgnoredException("色图正在处理!")
_ulmt.set_True(event.user_id)
@run_postprocessor
async def do_something(
matcher: Matcher,
exception: Optional[Exception],
bot: Bot,
event: Event,
state: T_State,
):
global setu_data_list, withdraw_message_id
if isinstance(event, MessageEvent):
if matcher.module == "send_setu":
# 解除占用
_ulmt.set_False(event.user_id)
tasks = []
# 撤回色图
for id_ in withdraw_message_id[:]:
tasks.append(asyncio.ensure_future(withdraw_message(bot, event, id_)))
withdraw_message_id.remove(id_)
await asyncio.gather(*tasks)
# 添加数据至数据库
try:
await add_data_to_database(setu_data_list)
logger.info("色图数据自动存储数据库成功...")
setu_data_list = []
except UninitializedError:
pass
setu = on_command(
"色图", aliases={"涩图", "不够色", "来一发", "再来点", "色图r"}, priority=5, block=True
)
setu_reg = on_regex("(.*)[份|发|张|个|次|点](.*)[瑟|色|涩]图", priority=5, block=True)
find_setu = on_command("查色图", priority=5, block=True)
@setu.handle()
async def _(bot: Bot, event: MessageEvent, state: T_State):
global withdraw_message_id
msg = get_message_text(event.json())
if isinstance(event, GroupMessageEvent):
impression = (
await SignGroupUser.ensure(event.user_id, event.group_id)
).impression
luox = get_luoxiang(impression)
if luox:
await setu.finish(luox)
if not _flmt.check(event.user_id):
await setu.finish("您冲得太快了,请稍候再冲", at_sender=True)
_flmt.start_cd(event.user_id)
r18 = 0
num = 1
# 是否看r18
if state["_prefix"]["raw_command"] == "色图r" and isinstance(
event, PrivateMessageEvent
):
r18 = 1
num = 10
elif state["_prefix"]["raw_command"] == "色图r" and isinstance(
event, GroupMessageEvent
):
await setu.finish(
random.choice(["这种不好意思的东西怎么可能给这么多人看啦", "羞羞脸!给我滚出克私聊!", "变态变态变态变态大变态!"])
)
# 有 数字 的话先尝试本地色图id
if msg and is_number(msg):
setu_list, code = await get_setu_list(int(msg), r18=r18)
if code != 200:
await setu.finish(setu_list[0], at_sender=True)
setu_img, code = await check_local_exists_or_download(setu_list[0])
msg_id = await setu.send(gen_message(setu_list[0]) + setu_img, at_sender=True)
logger.info(
f"(USER {event.user_id}, GROUP "
f"{event.group_id if isinstance(event, GroupMessageEvent) else 'private'})"
f" 发送色图 {setu_list[0].local_id}.png"
)
withdraw_message_id.append(msg_id["message_id"])
return
await send_setu_handle(setu, event, state["_prefix"]["raw_command"], msg, num, r18)
num_key = {
'一': 1,
'二': 2,
'两': 2,
'双': 2,
'三': 3,
'四': 4,
'五': 5,
'六': 6,
'七': 7,
'八': 8,
'九': 9
"一": 1,
"二": 2,
"两": 2,
"双": 2,
"三": 3,
"四": 4,
"五": 5,
"六": 6,
"七": 7,
"八": 8,
"九": 9,
}
@setu_reg.handle()
async def _(bot: Bot, event: MessageEvent, state: T_State):
if event.message_type == "group":
impression = (await SignGroupUser.ensure(event.user_id, event.group_id)).impression
if isinstance(event, GroupMessageEvent):
impression = (
await SignGroupUser.ensure(event.user_id, event.group_id)
).impression
luox = get_luoxiang(impression)
if luox:
await setu.finish(luox, at_sender=True)
if _ulmt.check(event.user_id):
await setu.finish(f"您有色图正在处理,请稍等")
_ulmt.set_True(event.user_id)
if not _flmt.check(event.user_id):
_ulmt.set_False(event.user_id)
await setu.finish('您冲得太快了,请稍候再冲', at_sender=True)
await setu.finish("您冲得太快了,请稍候再冲", at_sender=True)
_flmt.start_cd(event.user_id)
msg = get_message_text(event.json())
num = 1
msg = re.search(r'(.*)[份|发|张|个|次|点](.*)[瑟|涩|色]图', msg)
msg = re.search(r"(.*)[份发张个次点](.*)[瑟涩色]图", msg)
# 解析 tags 以及 num
if msg:
num = msg.group(1)
keyword = msg.group(2)
if keyword:
if keyword[-1] == '的':
keyword = keyword[:-1]
tags = msg.group(2)
if tags:
tags = tags[:-1] if tags[-1] == "的" else tags
if num:
num = num[-1]
if num_key.get(num):
@@ -260,111 +217,124 @@ async def _(bot: Bot, event: MessageEvent, state: T_State):
num = 1
else:
return
# try:
try:
setu_data = json.load(open(f'{TXT_PATH}/setu_data.json', 'r', encoding='utf8'))
except Exception:
setu_data = {}
if not keyword:
for _ in range(num):
try:
img, index = get_setu('', setu_data)
if not img:
break
await setu_reg.send(img)
except Exception as e:
await setu_reg.send('有图太色了发不出来...')
else:
logger.info(
f"USER {event.user_id} GROUP {event.group_id if event.message_type != 'private' else 'private'}"
f" 发送 {index} 色图成功")
_ulmt.set_False(event.user_id)
else:
_ulmt.set_False(event.user_id)
return
if keyword.find('真寻') != -1:
_ulmt.set_False(event.user_id)
await setu.finish('咳咳咳,虽然我很可爱,但是我木有自己的色图~~~有的话记得发我一份呀')
urls, text_list, code = await get_setu_urls(keyword, num)
if code == 200:
for i in range(len(urls)):
try:
setu_img, index = await search_online_setu(urls[i])
if index == -1:
setu_img, index = get_setu('', setu_data, keyword)
if setu_img:
await setu_reg.send(setu_img)
logger.error(f'色图下载失败 本地发送 {index}.jpg')
else:
await setu_reg.send(text_list[i] + '\n' + setu_img)
except ActionFailed as e:
await setu_reg.send('这图太色了,会教坏小孩子的,不给看..')
except Exception as e:
setu_img, index = get_setu('', setu_data, keyword)
if setu_img:
await setu_reg.send(setu_img)
logger.error(f'色图发送错误 e:{e}')
else:
logger.info(
f"USER {event.user_id} GROUP {event.group_id if event.message_type != 'private' else 'private'}"
f" 发送 {keyword} {num}连 色图成功")
else:
_ulmt.set_False(event.user_id)
await setu_reg.finish(urls, at_sender=True)
_ulmt.set_False(event.user_id)
await send_setu_handle(setu_reg, event, tags, num, 0)
@find_setu.args_parser
async def _(bot: Bot, event: MessageEvent, state: T_State):
if str(event.message) == '取消':
await find_setu.finish('取消了操作', at_sender=True)
if str(event.message) == "取消":
await find_setu.finish("取消了操作", at_sender=True)
imgs = get_message_imgs(event.json())
if not imgs:
await find_setu.reject("不搞错了,俺要图!")
state['img'] = imgs[0]
state["img"] = imgs[0]
@find_setu.handle()
async def _(bot: Bot, event: MessageEvent, state: T_State):
if get_message_text(event.json()) in ['帮助']:
await find_setu.finish('通过图片获取本地色图id\n\t示例:查色图(图片)')
if get_message_text(event.json()) in ["帮助"]:
await find_setu.finish("通过图片获取本地色图id\n\t示例:查色图(图片)")
imgs = get_message_imgs(event.json())
if imgs:
state['img'] = imgs[0]
state["img"] = imgs[0]
@find_setu.got('img', prompt="速速来图!")
@find_setu.got("img", prompt="速速来图!")
async def _(bot: Bot, event: MessageEvent, state: T_State):
img = state['img']
img = state["img"]
await find_setu.send(await find_img_index(img, event.user_id), at_sender=True)
@delete_setu.handle()
async def _(bot: Bot, event: MessageEvent, state: T_State):
_id = get_message_text(event.json())
if _id:
flag, text = delete_img(int(_id))
if flag:
await delete_setu.finish(f'删除色图 id:{_id} 成功', at_sender=True)
else:
await delete_setu.finish(f'删除色图 id:{_id} 失败,{text}', at_sender=True)
@upload_setu.handle()
async def _(bot: Bot, event: MessageEvent, state: T_State):
imgs = get_message_imgs(event.json())
if imgs:
lens, add_count = await add_img(imgs)
if add_count == 0:
await upload_setu.finish('上传的涩图已在色图库中,未成功上传,可以通过 查色图 来查看色图在图库中的id', at_sender=True)
id_s = ''
for i in range(add_count, 0, -1):
id_s += f'{lens - add_count} '
await upload_setu.finish(f"这次一共为 色图 库 添加了 {add_count} 张图片\n"
f"依次的Id为:{id_s[:-1]}\n"
f"小真寻感谢您对图库的扩充!WW", at_sender=True)
logger.info(
f"(USER {event.user_id}, GROUP {event.group_id if event.message_type != 'private' else 'private'})"
f" 上传色图ID:{id_s[:-1]}")
async def send_setu_handle(
matcher: Type[Matcher], event: MessageEvent, command: str, msg: str, num: int, r18: int
):
global setu_data_list, withdraw_message_id
error_info = ""
# 非 id,在线搜索
tags = msg.split()
# 真寻的色图?怎么可能
if "真寻" in tags:
await matcher.finish("咳咳咳,虽然我很可爱,但是我木有自己的色图~~~有的话记得发我一份呀")
# 本地先拿图,下载失败补上去
setu_list, code = await get_setu_list(tags=msg.split(), r18=r18)
if not ONLY_USE_LOCAL_SETU and (tags or not setu_list or num > len(setu_list)):
# 先尝试获取在线图片
urls, text_list, add_databases_list, code = await get_setu_urls(tags, num, r18, command)
for x in add_databases_list:
setu_data_list.append(x)
# 未找到符合的色图,想来本地应该也没有
if code == 401:
await setu.finish(urls[0], at_sender=True)
if code == 200:
for i in range(len(urls)):
try:
setu_img, index = await search_online_setu(urls[i])
# 下载成功的话
if index != -1:
logger.info(
f"(USER {event.user_id}, GROUP "
f"{event.group_id if isinstance(event, GroupMessageEvent) else 'private'})"
f" 发送色图 {index}.png"
)
msg_id = await matcher.send(Message(f"{text_list[i]}\n{setu_img}"))
withdraw_message_id.append(msg_id["message_id"])
else:
if setu_list:
setu_image = random.choice(setu_list)
setu_list.remove(setu_image)
msg_id = await matcher.send(
Message(
gen_message(setu_image)
+ (await check_local_exists_or_download(setu_image))[0]
)
)
logger.info(
f"(USER {event.user_id}, GROUP "
f"{event.group_id if isinstance(event, GroupMessageEvent) else 'private'})"
f" 发送本地色图 {setu_image.local_id}.png"
)
else:
msg_id = await matcher.send(text_list[i] + "\n" + setu_img)
withdraw_message_id.append(msg_id["message_id"])
except ActionFailed:
await matcher.finish("坏了,这张图色过头了,我自己看看就行了!", at_sender=True)
return
# 本地无图 或 超过上下限
if code != 200 or (not setu_list and ONLY_USE_LOCAL_SETU):
await matcher.finish(setu_list[0], at_sender=True)
elif not setu_list:
await matcher.finish(error_info, at_sender=True)
# 开始发图
for _ in range(num):
if not setu_list:
await setu.finish("坏了,已经没图了,被榨干了!")
setu_image = random.choice(setu_list)
setu_list.remove(setu_image)
try:
msg_id = await matcher.send(
Message(
gen_message(setu_image)
+ (await check_local_exists_or_download(setu_image))[0]
)
)
withdraw_message_id.append(msg_id["message_id"])
logger.info(
f"(USER {event.user_id}, GROUP "
f"{event.group_id if isinstance(event, GroupMessageEvent) else 'private'})"
f" 发送本地色图 {setu_image.local_id}.png"
)
except ActionFailed:
await matcher.finish("坏了,这张图色过头了,我自己看看就行了!", at_sender=True)
# 撤回图片
async def withdraw_message(bot: Bot, event: MessageEvent, id_: int):
if WITHDRAW_SETU_TIME[0]:
if (
(WITHDRAW_SETU_TIME[1] == 0 and isinstance(event, PrivateMessageEvent))
or (WITHDRAW_SETU_TIME[1] == 1 and isinstance(event, GroupMessageEvent))
or WITHDRAW_SETU_TIME[1] == 2
):
await asyncio.sleep(WITHDRAW_SETU_TIME[0])
await bot.delete_msg(message_id=id_, self_id=int(bot.self_id))
logger.info(f"自动撤回色图 消息id:{id_}")
-62
View File
@@ -1,62 +0,0 @@
import os
import imagehash
from PIL import Image
try:
import ujson as json
except ModuleNotFoundError:
import json
IMAGE_PATH = r"/home/hibiki/hibikibot/resources/img/"
TXT_PATH = r"/home/hibiki/hibikibot/resources/txt/"
def get_img_hash(image_file):
with open(image_file, 'rb') as fp:
hash_value = imagehash.average_hash(Image.open(fp))
return hash_value
def check_file_index():
for dir_name in ['setu/', 'r18/']:
lens = len(os.listdir(IMAGE_PATH + dir_name))
for i in range(lens):
if i > lens:
return
if not os.path.exists(f"{IMAGE_PATH}{dir_name}{i}.jpg"):
os.rename(f"{IMAGE_PATH}{dir_name}{lens}.jpg", f"{IMAGE_PATH}{dir_name}{i}.jpg")
print(f'{lens}.jpg --> {i}.jpg')
lens -= 1
if not i % 100:
print(f'已检测 {i} 份数据')
print(f'{dir_name} 检测完毕')
def check_setu_hash():
check_file_index()
for dir_name in ['setu/', 'r18/']:
img_data = {}
if dir_name == 'setu/':
fn = 'setu_img_hash.json'
else:
fn = 'r18_setu_img_hash.json'
file_list_len = len(os.listdir(IMAGE_PATH + dir_name)) - 1
print(file_list_len)
for i in range(file_list_len):
file = f"{i}.jpg"
index = file.split(".")[0]
img_hash = str(get_img_hash(IMAGE_PATH + dir_name + file))
print(f'{index}.jpg --> {img_hash}')
if img_hash in img_data.values():
k = [k for k, v in img_data.items() if v == img_hash]
print(f'文件 {index}.jpg 与 {k}.jpg 重复,使用 {file_list_len}.jpg 进行替换')
os.remove(IMAGE_PATH + dir_name + file)
os.rename(IMAGE_PATH + f'{dir_name}{file_list_len}.jpg', IMAGE_PATH + dir_name + file)
file_list_len -= 1
continue
img_data[index] = img_hash
# print(f'{index}.jpg --> {img_data}')
with open(TXT_PATH + fn, 'w') as f:
json.dump(img_data, f, indent=4)
check_setu_hash()
+171 -204
View File
@@ -1,260 +1,227 @@
from configs.path_config import IMAGE_PATH, TXT_PATH
import os
import random
from utils.init_result import image
from configs.config import LOLICON_KEY
from configs.path_config import IMAGE_PATH
from utils.message_builder import image
from services.log import logger
from aiohttp.client_exceptions import ClientConnectorError
from utils.image_utils import get_img_hash
from asyncpg.exceptions import UniqueViolationError
from utils.utils import get_local_proxy
from asyncio.exceptions import TimeoutError
from typing import List, Optional
from models.setu import Setu
import aiohttp
import aiofiles
from services.log import logger
from utils.img_utils import get_img_hash
from utils.utils import get_local_proxy, is_number
from asyncio.exceptions import TimeoutError
from models.count_user import UserCount
from configs.config import DOWNLOAD_SETU
import os
import random
try:
import ujson as json
except ModuleNotFoundError:
import json
url = "https://api.lolicon.app/setu/"
path = '_setu/'
url = "https://api.lolicon.app/setu/v2"
path = "_setu/"
r18_path = "_r18/"
async def get_setu_urls(keyword: str, num: int = 1, r18: int = 0):
# print(keyword)
if r18 == 1:
file_name = 'setu_r18_url.json'
else:
file_name = 'setu_url.json'
try:
with open(TXT_PATH + file_name, 'r', encoding='utf8') as f:
txt_data = json.load(f)
except (FileNotFoundError, ValueError):
txt_data = {}
txt_urls = [txt_data[x]['img_url'] for x in txt_data.keys()]
# 获取url
async def get_setu_urls(
tags: List[str], num: int = 1, r18: int = 0, command: str = ''
) -> "List[str], List[str], List[tuple], int":
tags = tags[:20] if len(tags) > 20 else tags
params = {
"apikey": LOLICON_KEY, # 添加apikey
'r18': r18, # 添加r18参数 0为否,1为是,2为混合
'keyword': keyword, # 若指定关键字,将会返回从插画标题、作者、标签中模糊搜索的结果
'num': 100, # 一次返回的结果数量,范围为1到10,不提供 APIKEY 时固定为1
'size1200': 1, # 是否使用 master_1200 缩略图,以节省流量或提升加载速度
"r18": r18, # 添加r18参数 0为否,1为是,2为混合
"tag": "|".join(tags), # 若指定tag
"num": 100, # 一次返回的结果数量,范围为1到10,不提供 APIKEY 时固定为1
"size1200": "original", # 是否使用 master_1200 缩略图,以节省流量或提升加载速度
}
urls = []
text_list = []
for count in range(3):
print(f'get_setu_url: count --> {count}')
async with aiohttp.ClientSession() as session:
add_databases_list = []
async with aiohttp.ClientSession() as session:
for count in range(3):
logger.info(f"get_setu_url: count --> {count}")
try:
async with session.get(url, proxy=get_local_proxy(), timeout=2, params=params) as response:
if response.status == 429:
return '调用达到上限,明日赶早呀~', '', 429
if response.status == 404:
return "网站裂开了...", '', 998
async with session.get(
url, proxy=get_local_proxy(), timeout=2, params=params
) as response:
if response.status == 200:
data = await response.json()
if data['code'] == 0:
lens = len(data['data'])
for i in range(lens):
img_url = data['data'][i]['url']
title = data['data'][i]['title']
author = data['data'][i]['author']
pid = data['data'][i]['pid']
if not data["error"]:
data = data["data"]
for i in range(len(data)):
img_url = data[i]["urls"]["original"]
img_url = (
img_url.replace("i.pixiv.cat", "i.pximg.net")
if "i.pixiv.cat" in img_url
else img_url
)
title = data[i]["title"]
author = data[i]["author"]
pid = data[i]["pid"]
urls.append(img_url)
text_list.append(f'title:{title}\nauthor:{author}\nPID:{pid}')
text_list.append(
f"title:{title}\nauthor:{author}\nPID:{pid}"
)
tags = []
for j in range(len(data['data'][i]['tags'])):
tags.append(data['data'][i]['tags'][j])
if img_url not in txt_urls:
save_setu_dict = {
'title': title,
'author': author,
'pid': pid,
'img_url': img_url,
'tags': tags
}
if str(pid) not in txt_data.keys():
txt_data[pid] = save_setu_dict
if DOWNLOAD_SETU:
with open(TXT_PATH + file_name, 'w', encoding='utf8') as f:
json.dump(txt_data, f, ensure_ascii=False, indent=4)
num = lens if num > lens else num
random_idx = random.sample(range(len(data['data'])), num)
for j in range(len(data[i]["tags"])):
tags.append(data[i]["tags"][j])
if command != '色图r':
if 'R-18' in tags:
tags.remove('R-18')
add_databases_list.append(
(
title,
author,
pid,
"",
img_url,
",".join(tags),
)
)
random_idx = random.sample(range(len(data)), num)
x_urls = []
x_text_lst = []
for x in random_idx:
x_urls.append(urls[x])
x_text_lst.append(text_list[x])
return x_urls, x_text_lst, 200
return x_urls, x_text_lst, add_databases_list, 200
else:
return "没找到符合条件的色图...", '', 401
except TimeoutError:
return ["没找到符合条件的色图..."], [], [], 401
except (TimeoutError, ClientConnectorError):
pass
return '我网线被人拔了..QAQ', '', 999
return ["我网线被人拔了..QAQ"], [], [], 999
async def search_online_setu(url: str):
async with aiohttp.ClientSession() as session:
headers = {
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.6;"
" rv:2.0.1) Gecko/20100101 Firefox/4.0.1",
"Referer": "https://www.pixiv.net",
}
async def search_online_setu(url_: str) -> "MessageSegment, int":
if "i.pixiv.cat" in url_:
url_ = url_.replace("i.pixiv.cat", "i.pximg.net")
async with aiohttp.ClientSession(headers=headers) as session:
for i in range(3):
print(f'search_online_setu --> {i}')
logger.info(f"search_online_setu --> {i}")
try:
async with session.get(url, proxy=get_local_proxy(), timeout=2) as res:
async with session.get(url_, proxy=get_local_proxy(), timeout=3) as res:
if res.status == 200:
index = str(random.randint(1, 100000))
async with aiofiles.open(IMAGE_PATH + 'temp/' + index + "_temp_setu.jpg", 'wb') as f:
async with aiofiles.open(
IMAGE_PATH + "temp/" + index + "_temp_setu.jpg", "wb"
) as f:
try:
await f.write(await res.read())
except TimeoutError:
# return '\n这图没下载过来~(网太差?)', -1, False
continue
logger.info(f"下载 lolicon图片 {url} 成功, id:{index}")
return image(f'{index}_temp_setu.jpg', 'temp'), index
logger.info(f"下载 lolicon图片 {url_} 成功, id:{index}")
return image(f"{index}_temp_setu.jpg", "temp"), index
else:
logger.warning(f"访问 lolicon图片 {url} 失败 status:{res.status}")
logger.warning(f"访问 lolicon图片 {url_} 失败 status:{res.status}")
# return '\n这图好难下载啊!QAQ', -1, False
except TimeoutError:
except (TimeoutError, ClientConnectorError):
pass
return '\n图片被小怪兽恰掉啦..!QAQ', -1
return "图片被小怪兽恰掉啦..!QAQ", -1
def get_setu(index: str, setu_data: dict, num: int = 1, tag: str = None, r18: bool = False):
if r18:
if not os.path.exists(IMAGE_PATH + '_r18'):
os.mkdir(IMAGE_PATH + '_r18')
length = len(os.listdir(IMAGE_PATH + '_r18'))
# 检测本地是否有id涩图,无的话则下载
async def check_local_exists_or_download(setu_image: Setu) -> "MessageSegment, int":
if setu_image.is_r18:
path_ = "_r18"
else:
if not os.path.exists(IMAGE_PATH + path):
os.mkdir(IMAGE_PATH + path)
length = len(os.listdir(IMAGE_PATH + path))
if length == 0:
return None, None
if not index:
index = random.randint(0, length - 1)
if tag:
all_tag_setu = [x for x in setu_data.keys() if tag in setu_data[x]['tags']]
if all_tag_setu:
index = random.choice(all_tag_setu)
if is_number(index):
if int(index) > length or int(index) < 0:
return f"超过当前上下限!({length - 1})", 999
else:
if setu_data:
index = str(index)
return f'id:{index}\n' \
f'title:{setu_data[index]["title"]}\n' \
f'author:{setu_data[index]["author"]}\n' \
f'PID:{setu_data[index]["pid"]}' + image(f'{index}.jpg', path), index
else:
return f'id:{index}' + image(f'{index}.jpg', path), index
return None, None
path_ = path
if os.path.exists(f"{IMAGE_PATH}/{path_}/{setu_image.local_id}.jpg"):
return image(f"{setu_image.local_id}.jpg", path_), 200
return await search_online_setu(setu_image.img_url)
# 添加涩图数据到数据库
async def add_data_to_database(lst: List[tuple]):
tmp = []
for x in lst:
if x not in tmp:
tmp.append(x)
if tmp:
for x in tmp:
try:
r18 = 1 if 'R-18' in x[5] else 0
idx = await Setu.get_image_count(r18)
await Setu.add_setu_data(
idx,
x[0],
x[1],
x[2],
x[3],
x[4],
x[5],
)
except UniqueViolationError:
pass
# 拿到本地色图列表
async def get_setu_list(
index: Optional[int] = None, tags: Optional[List[str]] = None, r18: int = 0
) -> "list, int":
if index:
image_count = await Setu.get_image_count(r18)
if index < 0 or index > image_count:
return [f"超过当前上下限!({image_count})"], 999
image_list = [await Setu.query_image(index, r18=r18)]
elif tags:
image_list = await Setu.query_image(tags=tags, r18=r18)
else:
image_list = await Setu.query_image(r18=r18)
if not image_list:
return ["没找到符合条件的色图..."], 998
return image_list, 200
# 初始化消息
def gen_message(setu_image: Setu, img_msg: bool = False):
local_id = setu_image.local_id
title = setu_image.title
author = setu_image.author
pid = setu_image.pid
return (
f"id:{local_id}\n"
f"title:{title}\n"
f"author:{author}\n"
f"PID:{pid}\n"
f"{image(f'{local_id}', f'{r18_path if setu_image.is_r18 else path}') if img_msg else ''}"
)
# 罗翔老师!
def get_luoxiang(impression):
probability = impression + 70
if probability < random.randint(1, 101):
return "我为什么要给你发这个?" + image(random.choice(os.listdir(IMAGE_PATH + "luoxiang/")), 'luoxiang') + \
"\n(快向小真寻签到提升好感度吧!)"
return (
"我为什么要给你发这个?"
+ image(random.choice(os.listdir(IMAGE_PATH + "luoxiang/")), "luoxiang")
+ "\n(快向小真寻签到提升好感度吧!)"
)
return None
async def check_r18_and_keyword(msg: str, user_id) -> 'str, int, int':
msg_list = msg.split(' ')
num = 1
r18 = 0
keyword = ''
if len(msg_list) == 1:
if msg_list[0].strip().lower() in ['r', 'r18']:
r18 = 1
num = 10
else:
keyword = msg_list[0]
elif len(msg_list) == 2:
keyword = msg_list[1].strip()
if msg_list[0].strip().lower() in ['r', 'r18']:
r18 = 1
num = 10
else:
keyword = msg[0]
if r18 == 1:
await UserCount.add_user(user_id)
return keyword, r18, num
async def find_img_index(img_url, user_id):
try:
setu_data = json.load(open(TXT_PATH + 'setu_data.json', encoding='utf8'))
except (FileNotFoundError, ValueError):
setu_data = {}
async with aiohttp.ClientSession() as session:
async with session.get(img_url, proxy=get_local_proxy(), timeout=5) as res:
async with aiofiles.open(IMAGE_PATH + f"temp/{user_id}_find_setu_index.jpg", 'wb') as f:
async with aiofiles.open(
IMAGE_PATH + f"temp/{user_id}_find_setu_index.jpg", "wb"
) as f:
await f.write(await res.read())
img_hash = str(get_img_hash(IMAGE_PATH + f"temp/{user_id}_find_setu_index.jpg"))
try:
index = str([setu_data[x]['img_hash'] for x in setu_data.keys()].index(img_hash))
return f"id:{index}\n" \
f"title:{setu_data[index]['title']}\n" \
f"author:{setu_data[index]['author']}\n" \
f"PID:{setu_data[index]['pid']}"
except ValueError:
return "该图不在色图库中!"
def delete_img(_id: int):
lens = len(os.listdir(IMAGE_PATH + path)) - 1
if _id < 0 or _id > lens:
return False, f'超过上下限限制,上限:{lens}'
try:
os.remove(IMAGE_PATH + path + f'{_id}.jpg')
if _id != lens:
setu_hash_dict = json.load(open(TXT_PATH + 'setu_data.json', encoding='utf8'))
setu_hash_dict[str(_id)] = setu_hash_dict[str(lens)]
os.rename(IMAGE_PATH + path + f'{lens}.jpg', IMAGE_PATH + path + f'{_id}.jpg')
with open(TXT_PATH + 'setu_data.json', 'w', encoding='utf8') as f:
json.dump(setu_hash_dict, f, ensure_ascii=False, indent=4)
return True, ''
except Exception as e:
logger.error(f'删除色图错误 e:{e}')
return False, str(type(e))
# 添加涩图
async def add_img(imgs: list):
index = 0
lens = len(os.listdir(IMAGE_PATH + path))
add_count = 0
setu_data_dict = json.load(open(TXT_PATH + 'setu_data.json', encoding='utf8'))
async with aiohttp.ClientSession() as session:
for img in imgs:
async with session.get(img, proxy=get_local_proxy(), timeout=5) as res:
async with aiofiles.open(IMAGE_PATH + f"temp/add_setu_check_{index}.jpg", 'wb') as f:
await f.write(await res.read())
index += 1
index -= 1
for i in range(index, -1, -1):
img_hash = str(get_img_hash(IMAGE_PATH + f"temp/add_setu_check_{index}.jpg"))
if img_hash not in [setu_data_dict[x]['img_hash'] for x in setu_data_dict.keys()]:
os.rename(IMAGE_PATH + f"temp/add_setu_check_{index}.jpg", IMAGE_PATH + path + f'/{lens}.jpg')
setu_data_dict[lens] = {
'title': 'not title',
'author': 'not author',
'pid': 'not pid',
'img_hash': img_hash,
'img_url': 'not url',
}
lens += 1
add_count += 1
if add_count:
with open(TXT_PATH + 'setu_data.json', 'w', encoding='utf8') as f:
json.dump(setu_data_dict, f, ensure_ascii=False, indent=4)
return lens, add_count
setu_img = await Setu.get_image_in_hash(img_hash)
if setu_img:
return (
f"id:{setu_img.local_id}\n"
f"title:{setu_img.title}\n"
f"author:{setu_img.author}\n"
f"PID:{setu_img.pid}"
)
return "该图不在色图库中或色图库未更新!"