Diuse/zhenxun_bot
1
0
Fork 0
mirror of https://github.com/zhenxun-org/zhenxun_bot.git synced 2026-10-03 02:30:01 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fit nonebot.beta2

Browse Source
This commit is contained in:
HibiKier
2022-02-19 18:20:19 +08:00
parent ab443a21ec
commit 7dd68cf1ef
212 changed files with 6857 additions and 5615 deletions
Download Patch File Download Diff File Expand all files Collapse all files
plugins/send_setu_/model.py → plugins/send_setu_/_model.py
Executable → Regular
+198 -198
View File
@@ -1,198 +1,198 @@
from services.db_context import db
from typing import List, Optional
class Setu(db.Model):
__tablename__ = "setu"
__table_args__ = {'extend_existing': True}
id = db.Column(db.Integer(), primary_key=True)
local_id = db.Column(db.Integer(), nullable=False)
title = db.Column(db.String(), nullable=False)
author = db.Column(db.String(), nullable=False)
pid = db.Column(db.BigInteger(), nullable=False)
img_hash = db.Column(db.String(), nullable=False)
img_url = db.Column(db.String(), nullable=False)
is_r18 = db.Column(db.Boolean(), nullable=False)
tags = db.Column(db.String())
_idx1 = db.Index("setu_pid_img_url_idx1", "pid", "img_url", unique=True)
@classmethod
async def add_setu_data(
cls,
local_id: int,
title: str,
author: str,
pid: int,
img_hash: str,
img_url: str,
tags: str,
):
"""
说明：
添加一份色图数据
参数：
:param local_id: 本地存储id
:param title: 标题
:param author: 作者
:param pid: 图片pid
:param img_hash: 图片hash值
:param img_url: 图片链接
:param tags: 图片标签
"""
if not await cls._check_exists(pid, img_url):
await cls.create(
local_id=local_id,
title=title,
author=author,
pid=pid,
img_hash=img_hash,
img_url=img_url,
is_r18=True if "R-18" in tags else False,
tags=tags,
)
@classmethod
async def query_image(
cls,
local_id: Optional[int] = None,
tags: Optional[List[str]] = None,
r18: int = 0,
limit: int = 50,
):
"""
说明：
通过tag查找色图
参数：
:param local_id: 本地色图 id
:param tags: tags
:param r18: 是否 r18，0：非r18 1：r18 2：混合
:param limit: 获取数量
"""
if local_id:
flag = True if r18 == 1 else False
return await cls.query.where(
(cls.local_id == local_id) & (cls.is_r18 == flag)
).gino.first()
if r18 == 0:
query = cls.query.where(cls.is_r18 == False)
elif r18 == 1:
query = cls.query.where(cls.is_r18 == True)
else:
query = cls.query
if tags:
for tag in tags:
query = query.where(cls.tags.contains(tag) | cls.title.contains(tag) | cls.author.contains(tag))
query = query.order_by(db.func.random()).limit(limit)
return await query.gino.all()
@classmethod
async def get_image_count(cls, r18: int = 0) -> int:
"""
说明：
查询图片数量
"""
flag = False if r18 == 0 else True
setattr(Setu, 'count', db.func.count(cls.local_id).label('count'))
count = await cls.select('count').where(cls.is_r18 == flag).gino.first()
return count[0]
@classmethod
async def get_image_in_hash(cls, img_hash: str) -> "Setu":
"""
说明：
通过图像hash获取图像信息
参数：
:param img_hash: = 图像hash值
"""
query = await cls.query.where(cls.img_hash == img_hash).gino.first()
return query
@classmethod
async def _check_exists(cls, pid: int, img_url: str) -> bool:
"""
说明：
检测图片是否存在
参数：
:param pid: 图片pid
:param img_url: 图片链接
"""
return bool(
await cls.query.where(
(cls.pid == pid) & (cls.img_url == img_url)
).gino.first()
)
@classmethod
async def delete_image(cls, pid: int) -> int:
"""
说明：
删除图片并替换
参数：
:param pid: 图片pid
"""
query = await cls.query.where(cls.pid == pid).gino.first()
if query:
is_r18 = query.is_r18
num = await cls.get_image_count(is_r18)
x = await cls.query.where((cls.is_r18 == is_r18) & (cls.local_id == num - 1)).gino.first()
_tmp_local_id = x.local_id
if x:
x.update(local_id=query.local_id).apply()
await cls.delete.where(cls.pid == pid).gino.status()
return _tmp_local_id
return -1
@classmethod
async def update_setu_data(
cls,
pid: int,
*,
local_id: Optional[int] = None,
title: Optional[str] = None,
author: Optional[str] = None,
img_hash: Optional[str] = None,
img_url: Optional[str] = None,
tags: Optional[str] = None,
) -> bool:
"""
说明：
根据PID修改图片数据
参数：
:param local_id: 本地id
:param pid: 图片pid
:param title: 标题
:param author: 作者
:param img_hash: 图片hash值
:param img_url: 图片链接
:param tags: 图片标签
"""
query = cls.query.where(cls.pid == pid).with_for_update()
image_list = await query.gino.all()
if image_list:
for image in image_list:
if local_id:
await image.update(local_id=local_id).apply()
if title:
await image.update(title=title).apply()
if author:
await image.update(author=author).apply()
if img_hash:
await image.update(img_hash=img_hash).apply()
if img_url:
await image.update(img_url=img_url).apply()
if tags:
await image.update(tags=tags).apply()
return True
return False
@classmethod
async def get_all_setu(cls) -> List["Setu"]:
"""
说明：
获取所有图片对象
"""
return await cls.query.gino.all()
from services.db_context import db
from typing import List, Optional
class Setu(db.Model):
__tablename__ = "setu"
__table_args__ = {'extend_existing': True}
id = db.Column(db.Integer(), primary_key=True)
local_id = db.Column(db.Integer(), nullable=False)
title = db.Column(db.String(), nullable=False)
author = db.Column(db.String(), nullable=False)
pid = db.Column(db.BigInteger(), nullable=False)
img_hash = db.Column(db.String(), nullable=False)
img_url = db.Column(db.String(), nullable=False)
is_r18 = db.Column(db.Boolean(), nullable=False)
tags = db.Column(db.String())
_idx1 = db.Index("setu_pid_img_url_idx1", "pid", "img_url", unique=True)
@classmethod
async def add_setu_data(
cls,
local_id: int,
title: str,
author: str,
pid: int,
img_hash: str,
img_url: str,
tags: str,
):
"""
说明：
添加一份色图数据
参数：
:param local_id: 本地存储id
:param title: 标题
:param author: 作者
:param pid: 图片pid
:param img_hash: 图片hash值
:param img_url: 图片链接
:param tags: 图片标签
"""
if not await cls._check_exists(pid, img_url):
await cls.create(
local_id=local_id,
title=title,
author=author,
pid=pid,
img_hash=img_hash,
img_url=img_url,
is_r18=True if "R-18" in tags else False,
tags=tags,
)
@classmethod
async def query_image(
cls,
local_id: Optional[int] = None,
tags: Optional[List[str]] = None,
r18: int = 0,
limit: int = 50,
):
"""
说明：
通过tag查找色图
参数：
:param local_id: 本地色图 id
:param tags: tags
:param r18: 是否 r18，0：非r18 1：r18 2：混合
:param limit: 获取数量
"""
if local_id:
flag = True if r18 == 1 else False
return await cls.query.where(
(cls.local_id == local_id) & (cls.is_r18 == flag)
).gino.first()
if r18 == 0:
query = cls.query.where(cls.is_r18 == False)
elif r18 == 1:
query = cls.query.where(cls.is_r18 == True)
else:
query = cls.query
if tags:
for tag in tags:
query = query.where(cls.tags.contains(tag) | cls.title.contains(tag) | cls.author.contains(tag))
query = query.order_by(db.func.random()).limit(limit)
return await query.gino.all()
@classmethod
async def get_image_count(cls, r18: int = 0) -> int:
"""
说明：
查询图片数量
"""
flag = False if r18 == 0 else True
setattr(Setu, 'count', db.func.count(cls.local_id).label('count'))
count = await cls.select('count').where(cls.is_r18 == flag).gino.first()
return count[0]
@classmethod
async def get_image_in_hash(cls, img_hash: str) -> "Setu":
"""
说明：
通过图像hash获取图像信息
参数：
:param img_hash: = 图像hash值
"""
query = await cls.query.where(cls.img_hash == img_hash).gino.first()
return query
@classmethod
async def _check_exists(cls, pid: int, img_url: str) -> bool:
"""
说明：
检测图片是否存在
参数：
:param pid: 图片pid
:param img_url: 图片链接
"""
return bool(
await cls.query.where(
(cls.pid == pid) & (cls.img_url == img_url)
).gino.first()
)
@classmethod
async def delete_image(cls, pid: int) -> int:
"""
说明：
删除图片并替换
参数：
:param pid: 图片pid
"""
query = await cls.query.where(cls.pid == pid).gino.first()
if query:
is_r18 = query.is_r18
num = await cls.get_image_count(is_r18)
x = await cls.query.where((cls.is_r18 == is_r18) & (cls.local_id == num - 1)).gino.first()
_tmp_local_id = x.local_id
if x:
x.update(local_id=query.local_id).apply()
await cls.delete.where(cls.pid == pid).gino.status()
return _tmp_local_id
return -1
@classmethod
async def update_setu_data(
cls,
pid: int,
*,
local_id: Optional[int] = None,
title: Optional[str] = None,
author: Optional[str] = None,
img_hash: Optional[str] = None,
img_url: Optional[str] = None,
tags: Optional[str] = None,
) -> bool:
"""
说明：
根据PID修改图片数据
参数：
:param local_id: 本地id
:param pid: 图片pid
:param title: 标题
:param author: 作者
:param img_hash: 图片hash值
:param img_url: 图片链接
:param tags: 图片标签
"""
query = cls.query.where(cls.pid == pid).with_for_update()
image_list = await query.gino.all()
if image_list:
for image in image_list:
if local_id:
await image.update(local_id=local_id).apply()
if title:
await image.update(title=title).apply()
if author:
await image.update(author=author).apply()
if img_hash:
await image.update(img_hash=img_hash).apply()
if img_url:
await image.update(img_url=img_url).apply()
if tags:
await image.update(tags=tags).apply()
return True
return False
@classmethod
async def get_all_setu(cls) -> List["Setu"]:
"""
说明：
获取所有图片对象
"""
return await cls.query.gino.all()
plugins/send_setu_/send_setu/__init__.py
+19 -60
View File
@@ -8,11 +8,10 @@ from typing import Optional, Type
from gino.exceptions import UninitializedError
from utils.utils import (
is_number,
get_message_text,
get_message_img,
)
from nonebot.typing import T_State
from nonebot.adapters.cqhttp import (
from nonebot.adapters.onebot.v11 import (
Bot,
MessageEvent,
GroupMessageEvent,
@@ -31,9 +30,11 @@ from .data_source import (
add_data_to_database,
get_setu_count,
)
from nonebot.adapters.cqhttp.exception import ActionFailed
from nonebot.adapters.onebot.v11.exception import ActionFailed
from configs.config import Config, NICKNAME
from utils.manager import withdraw_message_manager
from nonebot.params import CommandArg, Command
from typing import Tuple
import re
try:
@@ -87,30 +88,17 @@ __plugin_configs__ = {
"INITIAL_SETU_PROBABILITY": {
"value": 0.7,
"help": "初始色图概率，总概率 = 初始色图概率 + 好感度",
"default_value": 0.7
"default_value": 0.7,
},
"DOWNLOAD_SETU": {
"value": True,
"help": "是否存储下载的色图，使用本地色图可以加快图片发送速度",
"default_value": True
"default_value": True,
},
"TIMEOUT": {
"value": 10,
"help": "色图下载超时限制(秒)",
"default_value": 10
},
"SHOW_INFO": {
"value": True,
"help": "是否显示色图的基本信息，如PID等",
"default_value": True
}
"TIMEOUT": {"value": 10, "help": "色图下载超时限制(秒)", "default_value": 10},
"SHOW_INFO": {"value": True, "help": "是否显示色图的基本信息，如PID等", "default_value": True},
}
Config.add_plugin_config(
"pixiv",
"PIXIV_NGINX_URL",
"i.pixiv.re",
help_="Pixiv反向代理"
)
Config.add_plugin_config("pixiv", "PIXIV_NGINX_URL", "i.pixiv.re", help_="Pixiv反向代理")
setu_data_list = []
@@ -125,7 +113,7 @@ async def do_something(
):
global setu_data_list
if isinstance(event, MessageEvent):
if matcher.module == "send_setu":
if matcher.plugin_name == "send_setu":
# 添加数据至数据库
try:
await add_data_to_database(setu_data_list)
@@ -141,12 +129,12 @@ setu = on_command(
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):
msg = get_message_text(event.json())
async def _(
event: MessageEvent, cmd: Tuple[str, ...] = Command(), arg: Message = CommandArg()
):
msg = arg.extract_plain_text().strip()
if isinstance(event, GroupMessageEvent):
impression = (
await SignGroupUser.ensure(event.user_id, event.group_id)
@@ -157,14 +145,10 @@ async def _(bot: Bot, event: MessageEvent, state: T_State):
r18 = 0
num = 1
# 是否看r18
if state["_prefix"]["raw_command"] == "色图r" and isinstance(
event, PrivateMessageEvent
):
if cmd[0] == "色图r" and isinstance(event, PrivateMessageEvent):
r18 = 1
num = 10
elif state["_prefix"]["raw_command"] == "色图r" and isinstance(
event, GroupMessageEvent
):
elif cmd[0] == "色图r" and isinstance(event, GroupMessageEvent):
await setu.finish(
random.choice(["这种不好意思的东西怎么可能给这么多人看啦", "羞羞脸！给我滚出克私聊！", "变态变态变态变态大变态！"])
)
@@ -187,7 +171,7 @@ async def _(bot: Bot, event: MessageEvent, state: T_State):
Config.get_config("send_setu", "WITHDRAW_SETU_MESSAGE"),
)
return
await send_setu_handle(setu, event, state["_prefix"]["raw_command"], msg, num, r18)
await send_setu_handle(setu, event, cmd[0], msg, num, r18)
num_key = {
@@ -206,7 +190,7 @@ num_key = {
@setu_reg.handle()
async def _(bot: Bot, event: MessageEvent, state: T_State):
async def _(event: MessageEvent, arg: Message = CommandArg()):
if isinstance(event, GroupMessageEvent):
impression = (
await SignGroupUser.ensure(event.user_id, event.group_id)
@@ -214,7 +198,7 @@ async def _(bot: Bot, event: MessageEvent, state: T_State):
luox = get_luoxiang(impression)
if luox:
await setu.finish(luox, at_sender=True)
msg = get_message_text(event.json())
msg = arg.extract_plain_text().strip()
num = 1
msg = re.search(r"(.*)[份发张个次点](.*)[瑟涩色]图", msg)
# 解析 tags 以及 num
@@ -239,31 +223,6 @@ async def _(bot: Bot, event: MessageEvent, state: T_State):
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)
imgs = get_message_img(event.json())
if not imgs:
await find_setu.reject("不搞错了，俺要图！")
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示例：查色图(图片)")
imgs = get_message_img(event.json())
if imgs:
state["img"] = imgs[0]
@find_setu.got("img", prompt="速速来图！")
async def _(bot: Bot, event: MessageEvent, state: T_State):
img = state["img"]
await find_setu.send(await find_img_index(img, event.user_id), at_sender=True)
async def send_setu_handle(
matcher: Type[Matcher],
event: MessageEvent,
plugins/send_setu_/send_setu/data_source.py
+16 -19
View File
@@ -1,4 +1,4 @@
from configs.path_config import IMAGE_PATH
from configs.path_config import IMAGE_PATH, TEMP_PATH
from utils.message_builder import image
from services.log import logger
from utils.image_utils import get_img_hash, compressed_image
@@ -7,7 +7,7 @@ from asyncio.exceptions import TimeoutError
from typing import List, Optional
from configs.config import NICKNAME, Config
from utils.http_utils import AsyncHttpx
from ..model import Setu
from .._model import Setu
import asyncio
import os
import random
@@ -19,8 +19,8 @@ except ModuleNotFoundError:
url = "https://api.lolicon.app/setu/v2"
path = "_setu/"
r18_path = "_r18/"
path = "_setu"
r18_path = "_r18"
# 获取url
@@ -93,13 +93,12 @@ async def search_online_setu(
logger.info(f"search_online_setu --> {i}")
try:
index = random.randint(1, 100000) if id_ is None else id_
path_ = "temp" if not path_ else path_
file = f"{index}_temp_setu.jpg" if not path_ else f"{index}.jpg"
if not os.path.exists(f"{IMAGE_PATH}/{path_}"):
os.mkdir(f"{IMAGE_PATH}/{path_}")
path_ = IMAGE_PATH / path_ if path_ else TEMP_PATH
file_name = f"{index}_temp_setu.jpg" if not path_ == TEMP_PATH else f"{index}.jpg"
path_.mkdir(parents=True, exist_ok=True)
if not await AsyncHttpx.download_file(
url_,
f"{IMAGE_PATH}/{path_}/{file}",
path_ / file_name,
timeout=Config.get_config("send_setu", "TIMEOUT"),
):
continue
@@ -111,8 +110,8 @@ async def search_online_setu(
compressed_image(
f"{IMAGE_PATH}/{path_}/{index}.jpg",
)
logger.info(f"下载 lolicon图片 {url_} 成功， id：{index}")
return image(file, path_), index
logger.info(f"下载 lolicon 图片 {url_} 成功， id：{index}")
return image(abspath=path_ / file_name), index
except TimeoutError:
pass
except Exception as e:
@@ -126,11 +125,9 @@ async def check_local_exists_or_download(setu_image: Setu) -> "MessageSegment, i
id_ = None
if Config.get_config("send_setu", "DOWNLOAD_SETU"):
id_ = setu_image.local_id
if setu_image.is_r18:
path_ = "_r18"
else:
path_ = path
if os.path.exists(f"{IMAGE_PATH}/{path_}/{setu_image.local_id}.jpg"):
path_ = r18_path if setu_image.is_r18 else path
file = IMAGE_PATH / path_ / f"{setu_image.local_id}.jpg"
if file.exists():
return image(f"{setu_image.local_id}.jpg", path_), 200
return await search_online_setu(setu_image.img_url, id_, path_)
@@ -202,7 +199,7 @@ def get_luoxiang(impression):
if probability < random.randint(1, 101):
return (
"我为什么要给你发这个？"
+ image(random.choice(os.listdir(IMAGE_PATH + "luoxiang/")), "luoxiang")
+ image(random.choice(os.listdir(IMAGE_PATH / "luoxiang")), "luoxiang")
+ f"\n(快向{NICKNAME}签到提升好感度吧！)"
)
return None
@@ -219,11 +216,11 @@ async def get_setu_count(r18: int) -> int:
async def find_img_index(img_url, user_id):
if not await AsyncHttpx.download_file(
img_url,
IMAGE_PATH + f"temp/{user_id}_find_setu_index.jpg",
TEMP_PATH / f"{user_id}_find_setu_index.jpg",
timeout=Config.get_config("send_setu", "TIMEOUT"),
):
return "检索图片下载上失败..."
img_hash = str(get_img_hash(IMAGE_PATH + f"temp/{user_id}_find_setu_index.jpg"))
img_hash = str(get_img_hash(TEMP_PATH / f"{user_id}_find_setu_index.jpg"))
setu_img = await Setu.get_image_in_hash(img_hash)
if setu_img:
return (
plugins/send_setu_/update_setu/__init__.py
+1 -3
View File
@@ -1,8 +1,6 @@
from utils.utils import scheduler
from nonebot import on_command
from nonebot.permission import SUPERUSER
from nonebot.typing import T_State
from nonebot.adapters import Bot, Event
from nonebot.rule import to_me
from .data_source import update_setu_img
from configs.config import Config
@@ -29,7 +27,7 @@ update_setu = on_command(
@update_setu.handle()
async def _(bot: Bot, event: Event, state: T_State):
async def _():
if Config.get_config("send_setu", "DOWNLOAD_SETU"):
await update_setu.send("开始更新色图...", at_sender=True)
await update_setu_img(True)
plugins/send_setu_/update_setu/data_source.py
+14 -16
View File
@@ -4,11 +4,10 @@ from datetime import datetime
from utils.image_utils import compressed_image, get_img_hash
from utils.utils import get_bot
from PIL import UnidentifiedImageError
from ..model import Setu
from .._model import Setu
from asyncpg.exceptions import UniqueViolationError
from configs.config import Config
from utils.http_utils import AsyncHttpx
from pathlib import Path
from nonebot import Driver
import nonebot
import os
@@ -17,13 +16,13 @@ import shutil
driver: Driver = nonebot.get_driver()
_path = Path(IMAGE_PATH)
_path = IMAGE_PATH
# 替换旧色图数据，修复local_id一直是50的问题
@driver.on_startup
async def update_old_setu_data():
path = Path(TEXT_PATH)
path = TEXT_PATH
setu_data_file = path / "setu_data.json"
r18_data_file = path / "r18_setu_data.json"
if setu_data_file.exists() or r18_data_file.exists():
@@ -79,9 +78,9 @@ async def update_old_setu_data():
# 删除色图rar文件夹
shutil.rmtree(Path(IMAGE_PATH) / "setu_rar", ignore_errors=True)
shutil.rmtree(Path(IMAGE_PATH) / "r18_rar", ignore_errors=True)
shutil.rmtree(Path(IMAGE_PATH) / "rar", ignore_errors=True)
shutil.rmtree(IMAGE_PATH / "setu_rar", ignore_errors=True)
shutil.rmtree(IMAGE_PATH / "r18_rar", ignore_errors=True)
shutil.rmtree(IMAGE_PATH / "rar", ignore_errors=True)
headers = {
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.6;"
@@ -104,12 +103,11 @@ async def update_setu_img(flag: bool = False):
for image in image_list:
count += 1
path = _path / "_r18" if image.is_r18 else _path / "_setu"
rar_path = Path(TEMP_PATH)
local_image = path / f"{image.local_id}.jpg"
path.mkdir(exist_ok=True, parents=True)
rar_path.mkdir(exist_ok=True, parents=True)
TEMP_PATH.mkdir(exist_ok=True, parents=True)
if not local_image.exists() or not image.img_hash:
temp_file = rar_path / f"{image.local_id}.jpg"
temp_file = TEMP_PATH / f"{image.local_id}.jpg"
if temp_file.exists():
temp_file.unlink()
url_ = image.img_url
@@ -120,29 +118,29 @@ async def update_setu_img(flag: bool = False):
)
try:
if not await AsyncHttpx.download_file(
url_, rar_path / f"{image.local_id}.jpg"
url_, TEMP_PATH / f"{image.local_id}.jpg"
):
continue
_success += 1
try:
if (
os.path.getsize(
rar_path / f"{image.local_id}.jpg",
TEMP_PATH / f"{image.local_id}.jpg",
)
> 1024 * 1024 * 1.5
):
compressed_image(
rar_path / f"{image.local_id}.jpg",
TEMP_PATH / f"{image.local_id}.jpg",
path / f"{image.local_id}.jpg",
)
else:
logger.info(
f"不需要压缩，移动图片{rar_path}/{image.local_id}.jpg "
f"不需要压缩，移动图片{TEMP_PATH}/{image.local_id}.jpg "
f"--> /{path}/{image.local_id}.jpg"
)
os.rename(
f"{rar_path}/{image.local_id}.jpg",
f"{path}/{image.local_id}.jpg",
TEMP_PATH / f"/{image.local_id}.jpg",
path / f"{image.local_id}.jpg",
)
except FileNotFoundError:
logger.warning(f"文件 {image.local_id}.jpg 不存在，跳过...")