mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-09-29 08:39:59 +08:00
refactor code
This commit is contained in:
hibiki
+171
-204
@@ -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 "该图不在色图库中或色图库未更新！"
|
||||
|
||||
Reference in New Issue
Block a user