Skip to content

Commit

Permalink
feat: add password config option
Browse files Browse the repository at this point in the history
  • Loading branch information
@akynazh
akynazh committed Feb 2, 2024
1 parent e7499ef commit a070b4f
Show file tree
Hide file tree
Showing 4 changed files with 67 additions and 31 deletions.
2 changes: 2 additions & 0 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,8 @@ use_cache:
redis_host:
# optional, your redis port (required if use_cache == 1)
redis_port:
# optional, your redis password
redis_password:
# required, enable nsfw or not, 1 yes | 0 no
enable_nsfw: 0
```
Expand Down
84 changes: 56 additions & 28 deletions bot.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,8 @@
apihelper.proxy = BOT_CFG.proxy_json
BOT = telebot.TeleBot(BOT_CFG.tg_bot_token)
BOT_DB = BotFileDb(PATH_RECORD_FILE)
BOT_CACHE_DB = BotCacheDb(host=BOT_CFG.redis_host, port=BOT_CFG.redis_port, use_cache=BOT_CFG.use_cache)
BOT_CACHE_DB = BotCacheDb(host=BOT_CFG.redis_host, port=BOT_CFG.redis_port,
password=BOT_CFG.redis_password, use_cache=BOT_CFG.use_cache)
BASE_UTIL = jvav.BaseUtil(BOT_CFG.proxy_addr)
DMM_UTIL = jvav.DmmUtil(BOT_CFG.proxy_addr_dmm)
JAVBUS_UTIL = jvav.JavBusUtil(BOT_CFG.proxy_addr)
Expand Down Expand Up @@ -289,7 +290,8 @@ def get_page_elements(
to_next = page_count
else:
to_next = page + 1
btn_to_first = InlineKeyboardButton(text="<<", callback_data=f"1:{key_type}")
btn_to_first = InlineKeyboardButton(
text="<<", callback_data=f"1:{key_type}")
btn_to_previous = InlineKeyboardButton(
text="<", callback_data=f"{to_previous}:{key_type}"
)
Expand Down Expand Up @@ -344,7 +346,8 @@ def get_stars_record(self, page=1):
def get_star_detail_record_by_name_id(self, star_name: str, star_id: str):
record, is_stars_exists, is_avs_exists = BOT_DB.check_has_record()
if not record:
self.send_msg(reason="尚无该演员收藏记录", op=f"获取演员 <code>{star_name}</code> 的更多信息")
self.send_msg(reason="尚无该演员收藏记录",
op=f"获取演员 <code>{star_name}</code> 的更多信息")
return
avs = []
star_avs = []
Expand Down Expand Up @@ -593,7 +596,8 @@ def get_av_by_id(
key=id, value=av, type=BotCacheDb.TYPE_AV, expire=3600 * 24 * 1
)
else:
BOT_CACHE_DB.set_cache(key=id, value=av, type=BotCacheDb.TYPE_AV)
BOT_CACHE_DB.set_cache(
key=id, value=av, type=BotCacheDb.TYPE_AV)
else:
av_score = av["score"]
is_cache = True
Expand Down Expand Up @@ -812,23 +816,27 @@ def get_sample_by_id(self, id: str):
code, samples = JAVBUS_UTIL.get_samples_by_id(id)
if not self.check_success(code, op_get_sample):
return
BOT_CACHE_DB.set_cache(key=id, value=samples, type=BotCacheDb.TYPE_SAMPLE)
BOT_CACHE_DB.set_cache(key=id, value=samples,
type=BotCacheDb.TYPE_SAMPLE)
# 发送图片列表
samples_imp = []
sample_error = False
for sample in samples:
samples_imp.append(InputMediaPhoto(sample))
if len(samples_imp) == 10: # 图片数目达到 10 张则发送一次
try:
BOT.send_media_group(chat_id=BOT_CFG.tg_chat_id, media=samples_imp)
BOT.send_media_group(
chat_id=BOT_CFG.tg_chat_id, media=samples_imp)
samples_imp = []
except Exception:
sample_error = True
self.send_msg_fail_reason_op(reason="图片解析失败", op=op_get_sample)
self.send_msg_fail_reason_op(
reason="图片解析失败", op=op_get_sample)
break
if samples_imp != [] and not sample_error:
try:
BOT.send_media_group(chat_id=BOT_CFG.tg_chat_id, media=samples_imp)
BOT.send_media_group(
chat_id=BOT_CFG.tg_chat_id, media=samples_imp)
except Exception:
self.send_msg_fail_reason_op(reason="图片解析失败", op=op_get_sample)

Expand Down Expand Up @@ -869,7 +877,8 @@ def watch_av_by_id(self, id: str, type: str):
from_site = "avgle"
pv_src = pv_avgle
pv_cache = {"from_site": from_site, "src": pv_src}
BOT_CACHE_DB.set_cache(key=id, value=pv_cache, type=BotCacheDb.TYPE_PV)
BOT_CACHE_DB.set_cache(
key=id, value=pv_cache, type=BotCacheDb.TYPE_PV)
else:
from_site = pv["from_site"]
pv_src = pv["src"]
Expand Down Expand Up @@ -907,7 +916,8 @@ def watch_av_by_id(self, id: str, type: str):
if code != 200:
self.send_msg(f"MissAv 视频地址: {BASE_URL_MISS_AV}/{id}")
return
BOT_CACHE_DB.set_cache(key=id, value=video, type=BotCacheDb.TYPE_FV)
BOT_CACHE_DB.set_cache(
key=id, value=video, type=BotCacheDb.TYPE_FV)
self.send_msg(
f"""MissAv 视频地址: {BASE_URL_MISS_AV}/{id}
Expand All @@ -926,7 +936,8 @@ def search_star_by_name(self, star_name: str) -> bool:
code, star = JAVBUS_UTIL.check_star_exists(star_name)
if not self.check_success(code, op_search_star):
return
BOT_CACHE_DB.set_cache(key=star_name, value=star, type=BotCacheDb.TYPE_STAR)
BOT_CACHE_DB.set_cache(
key=star_name, value=star, type=BotCacheDb.TYPE_STAR)
if star_name_origin != star_name:
BOT_CACHE_DB.set_cache(
key=star_name_origin,
Expand All @@ -936,7 +947,8 @@ def search_star_by_name(self, star_name: str) -> bool:
star_id = star["star_id"]
star_name = star["star_name"]
if BOT_DB.check_star_exists_by_id(star_id=star_id):
self.get_star_detail_record_by_name_id(star_name=star_name, star_id=star_id)
self.get_star_detail_record_by_name_id(
star_name=star_name, star_id=star_id)
return True
markup = InlineKeyboardMarkup()
markup.row(
Expand Down Expand Up @@ -974,9 +986,11 @@ def get_top_stars(self, page=1):
code, stars = DMM_UTIL.get_top_stars(page)
if not self.check_success(code, op_get_top_stars):
return
BOT_CACHE_DB.set_cache(key=page, value=stars, type=BotCacheDb.TYPE_RANK)
BOT_CACHE_DB.set_cache(key=page, value=stars,
type=BotCacheDb.TYPE_RANK)
stars_tmp = [None] * 80
stars = stars_tmp[: ((page - 1) * 20)] + stars + stars_tmp[((page - 1) * 20):]
stars = stars_tmp[: ((page - 1) * 20)] + stars + \
stars_tmp[((page - 1) * 20):]
col, row = 4, 5
objs, page_btns, title = self.get_page_elements(
objs=stars, page=page, col=4, row=5, key_type=BotKey.KEY_GET_TOP_STARS
Expand Down Expand Up @@ -1016,7 +1030,8 @@ def get_more_magnets_by_id(self, id: str):
if not av:
return
magnets = av["magnets"]
BOT_CACHE_DB.set_cache(key=id, value=magnets, type=BotCacheDb.TYPE_MAGNET)
BOT_CACHE_DB.set_cache(key=id, value=magnets,
type=BotCacheDb.TYPE_MAGNET)
msg = ""
for magnet in magnets:
magnet_tags = ""
Expand All @@ -1041,7 +1056,8 @@ def get_more_magnets_by_id(self, id: str):

def get_star_new_avs_by_name_id(self, star_name: str, star_id: str):
op_get_star_new_avs = f"获取 <code>{star_name}</code> 最新 av"
ids = BOT_CACHE_DB.get_cache(key=star_id, type=BotCacheDb.TYPE_NEW_AVS_OF_STAR)
ids = BOT_CACHE_DB.get_cache(
key=star_id, type=BotCacheDb.TYPE_NEW_AVS_OF_STAR)
if not ids:
code, ids = JAVBUS_UTIL.get_new_ids_by_star_id(star_id=star_id)
if not self.check_success(code, op_get_star_new_avs):
Expand Down Expand Up @@ -1102,7 +1118,8 @@ def send_bts(self, q, bts):

def random_get_new_av(self):
page = random.randint(1, JAVLIB_UTIL.MAX_RANK_PAGE)
ids = BOT_CACHE_DB.get_cache(key=page, type=BotCacheDb.TYPE_JLIB_PAGE_NEW_AVS)
ids = BOT_CACHE_DB.get_cache(
key=page, type=BotCacheDb.TYPE_JLIB_PAGE_NEW_AVS)
if not ids:
code, ids = JAVLIB_UTIL.get_random_ids_from_rank_by_page(
page=page, list_type=1
Expand All @@ -1119,7 +1136,8 @@ def random_get_new_av(self):

def random_get_nice_av(self):
page = random.randint(1, JAVLIB_UTIL.MAX_RANK_PAGE)
ids = BOT_CACHE_DB.get_cache(key=page, type=BotCacheDb.TYPE_JLIB_PAGE_NICE_AVS)
ids = BOT_CACHE_DB.get_cache(
key=page, type=BotCacheDb.TYPE_JLIB_PAGE_NICE_AVS)
if not ids:
code, ids = JAVLIB_UTIL.get_random_ids_from_rank_by_page(
page=page, list_type=0
Expand All @@ -1140,7 +1158,8 @@ def random_get_nice_star_avs(self, star_name_ori):
)
if not avs:
star_name_ja = self.get_star_ja_name_by_zh_name(star_name_ori)
code, avs = DMM_UTIL.get_nice_avs_by_star_name(star_name=star_name_ja)
code, avs = DMM_UTIL.get_nice_avs_by_star_name(
star_name=star_name_ja)
if self.check_success(code, f"获取演员 {star_name_ori} 的高分 av"):
avs = avs[:60]
BOT_CACHE_DB.set_cache(
Expand Down Expand Up @@ -1191,7 +1210,8 @@ def handle_callback(call):
tmp = content.split("|")
star_name = tmp[0]
star_id = tmp[1]
bot_utils.get_star_new_avs_by_name_id(star_name=star_name, star_id=star_id)
bot_utils.get_star_new_avs_by_name_id(
star_name=star_name, star_id=star_id)
elif key_type == BotKey.KEY_RECORD_STAR_BY_STAR_NAME_ID:
s = content.find("|")
star_name = content[:s]
Expand Down Expand Up @@ -1238,14 +1258,16 @@ def handle_callback(call):
if BOT_DB.undo_record_id(id=content):
bot_utils.send_msg_success_op(op_undo_record_av)
else:
bot_utils.send_msg_fail_reason_op(reason="文件解析出错", op=op_undo_record_av)
bot_utils.send_msg_fail_reason_op(
reason="文件解析出错", op=op_undo_record_av)
elif key_type == BotKey.KEY_UNDO_RECORD_STAR_BY_STAR_NAME_ID:
s = content.find("|")
op_undo_record_star = f"取消收藏演员 <code>{content[:s]}</code>"
if BOT_DB.undo_record_star_by_id(star_id=content[s + 1:]):
bot_utils.send_msg_success_op(op_undo_record_star)
else:
bot_utils.send_msg_fail_reason_op(reason="文件解析出错", op=op_undo_record_star)
bot_utils.send_msg_fail_reason_op(
reason="文件解析出错", op=op_undo_record_star)
elif key_type == BotKey.KEY_SEARCH_STAR_BY_NAME:
star_name = content
star_name_alias = ""
Expand Down Expand Up @@ -1305,7 +1327,8 @@ def handle_message(message):
bot_utils.send_msg_fail_reason_op(reason="尚无收藏记录", op="获取收藏记录文件")
return
BOT.send_document(
chat_id=BOT_CFG.tg_chat_id, document=types.InputFile(PATH_RECORD_FILE)
chat_id=BOT_CFG.tg_chat_id, document=types.InputFile(
PATH_RECORD_FILE)
)
elif msg_cmd == "/rank":
bot_utils.get_top_stars(1)
Expand All @@ -1328,18 +1351,22 @@ def handle_message(message):
bot_utils.send_msg_fail_reason_op(reason="未找到结果", op=op)
return
if BOT_CFG.enable_nsfw == "0":
bts = list(filter(lambda bt: bt['category'] != "nsfw", bts))
bts = list(filter(lambda bt: "nsfw" not in bt['name'].lower(), bts))
bts = list(
filter(lambda bt: bt['category'] != "nsfw", bts))
bts = list(
filter(lambda bt: "nsfw" not in bt['name'].lower(), bts))
bts = bts[:5]
BOT_CACHE_DB.set_cache(key=msg, value=bts, type=BotCacheDb.TYPE_BT)
BOT_CACHE_DB.set_cache(
key=msg, value=bts, type=BotCacheDb.TYPE_BT)
bot_utils.send_bts(msg, bts)
else:
ids = [id.lower() for id in ids]
ids = set(ids)
ids_msg = ", ".join(ids)
bot_utils.send_msg(f"检测到番号: {ids_msg}, 开始搜索......")
for i, id in enumerate(ids):
threading.Thread(target=bot_utils.get_av_by_id, args=(id,)).start()
threading.Thread(target=bot_utils.get_av_by_id,
args=(id,)).start()


@BOT.callback_query_handler(func=lambda call: True)
Expand Down Expand Up @@ -1370,7 +1397,8 @@ def main():
except Exception as e:
LOG.error(f"无法连接到机器人: {e}")
return
BOT.set_my_commands([types.BotCommand(cmd, BOT_CMDS[cmd]) for cmd in BOT_CMDS])
BOT.set_my_commands([types.BotCommand(cmd, BOT_CMDS[cmd])
for cmd in BOT_CMDS])
BOT.infinity_polling()


Expand Down
1 change: 1 addition & 0 deletions config.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ def __init__(self, path_config_file: str):
self.use_cache = str(config["use_cache"]) if config["use_cache"] else "0"
self.redis_host = str(config["redis_host"]) if config["redis_host"] else ""
self.redis_port = str(config["redis_port"]) if config["redis_port"] else ""
self.redis_password = str(config["redis_password"]) if config["redis_password"] else ""
self.enable_nsfw = str(config["enable_nsfw"]) if config["enable_nsfw"] else "0"
# set
self.proxy_addr_dmm = ""
Expand Down
11 changes: 8 additions & 3 deletions database.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -235,7 +235,8 @@ class BotCacheDb:
"prefix": "jlib-page-new-avs-",
"expire": 3600 * 24 * 2,
}
CACHE_STAR_JA_NAME = {"prefix": "star-ja-name-", "expire": 3600 * 24 * 30 * 6}
CACHE_STAR_JA_NAME = {"prefix": "star-ja-name-",
"expire": 3600 * 24 * 30 * 6}
CACHE_NEW_AVS_OF_STAR = {
"prefix": "new-avs-of-star-",
"expire": 3600 * 24 * 12,
Expand Down Expand Up @@ -275,18 +276,22 @@ class BotCacheDb:
TYPE_BT: CACHE_BT,
}

def __init__(self, host: str, port: int, use_cache: str):
def __init__(self, host: str, port: int, password: str, use_cache: str):
"""初始化
:param str host: ip 地址
:param int port: 端口
:param str password: 密码
:param str use_cache: 是否使用缓存
"""
self.use_cache = use_cache
self.cache = None
if self.use_cache == "1":
try:
self.cache = redis.Redis(host=host, port=port)
if password:
self.cache = redis.Redis(host=host, port=port, password=password)
else:
self.cache = redis.Redis(host=host, port=port)
self.cache.ping()
LOG.info(f"连接到 redis 服务: {host}:{port}")
except Exception as e:
Expand Down

0 comments on commit a070b4f

Please sign in to comment.