xtex/akari-bot
Archived
1
0
Fork 0
Code Issues Pull requests Activity

update

Browse source
This commit is contained in:
yzhh 2024-01-11 18:25:49 +08:00
parent 7dac27ccff
commit 1660118cc6
6 changed files with 130 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
core/utils/http.py
View file

@ -39,7 +39,7 @@ def private_ip_check(url: str):
async def get_url(url: str, status_code: int = False, headers: dict = None, params: dict = None, fmt=None, timeout=20,
attempt=3,
request_private_ip=False, logging_err_resp=True):
request_private_ip=False, logging_err_resp=True, cookies=None):
"""利用AioHttp获取指定url的内容。
:param url: 需要获取的url。
@ -51,6 +51,7 @@ async def get_url(url: str, status_code: int = False, headers: dict = None, para
:param attempt: 指定请求尝试次数。
:param request_private_ip: 是否允许请求私有IP。
:param logging_err_resp: 是否记录错误响应。
:param cookies: 使用的 cookies
:returns: 指定url的内容（字符串）。
"""
@ -63,6 +64,8 @@ async def get_url(url: str, status_code: int = False, headers: dict = None, para
async with aiohttp.ClientSession(headers=headers,
connector=TCPConnector(verify_ssl=False) if debug else None, ) as session:
if cookies:
session.cookie_jar.update_cookies(cookies)
try:
async with session.get(url, timeout=aiohttp.ClientTimeout(total=timeout), headers=headers,
proxy=proxy, params=params) as req:
@ -92,7 +95,7 @@ async def get_url(url: str, status_code: int = False, headers: dict = None, para
async def post_url(url: str, data: any = None, status_code: int = False, headers: dict = None, fmt=None, timeout=20,
attempt=3, request_private_ip=False, logging_err_resp=True):
attempt=3, request_private_ip=False, logging_err_resp=True, cookies=None):
'''发送POST请求。
:param url: 需要发送的url。
:param data: 需要发送的数据。
@ -113,6 +116,8 @@ async def post_url(url: str, data: any = None, status_code: int = False, headers
async with aiohttp.ClientSession(headers=headers,
connector=TCPConnector(verify_ssl=False) if debug else None, ) as session:
if cookies:
session.cookie_jar.update_cookies(cookies)
try:
async with session.post(url, data=data, headers=headers,
timeout=aiohttp.ClientTimeout(total=timeout),

40
modules/wiki/audit.py
View file

@ -1,11 +1,15 @@
from datetime import datetime
from datetime import datetime, timedelta
from config import Config
from core.builtins import Bot, Plain, Image
from core.component import module
from core.scheduler import DateTrigger
from core.logger import Logger
from core.utils.image_table import image_table_render, ImageTable
from modules.wiki.utils.bot import BotAccount, LoginFailed
from modules.wiki.utils.dbutils import Audit
from modules.wiki.utils.wikilib import WikiLib
from modules.wiki.utils.dbutils import BotAccount as BotAccountDB
aud = module('wiki_audit', required_superuser=True,
@ -117,3 +121,37 @@ async def _(msg: Bot.MessageSession):
for bl in block_list:
wikis.append(f'{bl[0]} ({bl[1]})')
await msg.finish('\n'.join(wikis))
@aud.handle('bot add <apiLink> <account> <password>')
async def _(msg: Bot.MessageSession):
api_link = msg.parsed_msg['<apiLink>']
account = msg.parsed_msg['<account>']
password = msg.parsed_msg['<password>']
check = await WikiLib(api_link).check_wiki_available()
if check.available:
try:
await BotAccount._login(api_link, account, password)
except LoginFailed as e:
Logger.error(f'Login failed: {e}')
await msg.finish(f'Login failed: {e}')
else:
await msg.finish('Login success')
BotAccountDB.add(api_link, account, password)
else:
result = msg.locale.t('wiki.message.error.query') + \
('\n' + msg.locale.t('wiki.message.error.info') + check.message if check.message != '' else '')
await msg.finish(result)
@aud.handle('bot remove <apiLink>')
async def _(msg: Bot.MessageSession):
api_link = msg.parsed_msg['<apiLink>']
BotAccountDB.remove(api_link)
await msg.finish('Done')
@aud.handle(DateTrigger(datetime.now() + timedelta(seconds=10)))
async def login_bots():
Logger.info('Start login wiki bot account...')
await BotAccount.login()

43
modules/wiki/utils/bot.py Normal file
View file

@ -0,0 +1,43 @@
import aiohttp
from core.logger import Logger
from .dbutils import BotAccount as BotAccountDB
class LoginFailed(Exception):
pass
class BotAccount:
cookies = {}
@staticmethod
async def _login(api_link, account, password):
lgtoken_url = f'{api_link}?action=query&meta=tokens&type=login&format=json'
PARAMS_1 = {
'action': "login",
'lgname': account,
'lgpassword': password,
'format': "json"
}
async with aiohttp.ClientSession() as session:
async with session.get(lgtoken_url) as req:
if req.status != 200:
raise LoginFailed(f'Login failed: {await req.text()}')
PARAMS_1['lgtoken'] = (await req.json())['query']['tokens']['logintoken']
async with session.post(api_link, data=PARAMS_1) as req:
if req.status != 200:
raise LoginFailed(f'Login failed: {await req.text()}')
Logger.info(f'Logged in to {api_link} as {account}')
return req.cookies
@classmethod
async def login(cls):
accounts = BotAccountDB.get_all()
for account in accounts:
try:
cls.cookies[account.apiLink] = await BotAccount._login(account.apiLink,
account.botAccount,
account.botPassword)
except LoginFailed as e:
Logger.error(f'Login failed: {e}')

27
modules/wiki/utils/dbutils.py
View file

@ -8,7 +8,7 @@ from tenacity import retry, stop_after_attempt
from core.builtins import MessageSession
from database import session, auto_rollback_error
from modules.wiki.utils.orm import WikiTargetSetInfo, WikiInfo, WikiAllowList, WikiBlockList
from modules.wiki.utils.orm import WikiTargetSetInfo, WikiInfo, WikiAllowList, WikiBlockList, WikiBotAccountList
class WikiTargetInfo:
@ -208,3 +208,28 @@ class Audit:
@auto_rollback_error
def get_block_list() -> list:
return session.query(WikiBlockList.apiLink, WikiBlockList.operator)
class BotAccount:
@staticmethod
@retry(stop=stop_after_attempt(3), reraise=True)
@auto_rollback_error
def add(api_link, bot_account, bot_password):
session.add_all([WikiBotAccountList(apiLink=api_link, botAccount=bot_account, botPassword=bot_password)])
session.commit()
session.expire_all()
return True
@staticmethod
@retry(stop=stop_after_attempt(3), reraise=True)
@auto_rollback_error
def remove(api_link):
session.delete(session.query(WikiBotAccountList).filter_by(apiLink=api_link).first())
session.commit()
session.expire_all()
return True
@staticmethod
@retry(stop=stop_after_attempt(3), reraise=True)
def get_all():
return session.query(WikiBotAccountList).all()

8
modules/wiki/utils/orm.py
View file

@ -42,4 +42,12 @@ class WikiBlockList(Base):
operator = Column(String(512))
class WikiBotAccountList(Base):
__tablename__ = table_prefix + 'WikiBotAccountList'
__table_args__ = {'extend_existing': True}
apiLink = Column(String(512), primary_key=True)
botAccount = Column(String(512))
botPassword = Column(String(512))
Session.create()

8
modules/wiki/utils/wikilib.py
View file

@ -16,6 +16,7 @@ from core.utils.http import get_url
from core.utils.i18n import Locale, default_locale
from core.exceptions import NoReportException
from modules.wiki.utils.dbutils import WikiSiteInfo as DBSiteInfo, Audit
from modules.wiki.utils.bot import BotAccount
web_render = CFG.get_url('web_render')
web_render_local = CFG.get_url('web_render_local')
@ -170,8 +171,13 @@ class WikiLib:
api = (web_render_local if use_local else web_render) + 'source?url=' + urllib.parse.quote(api)
request_local = True
break
cookies = None
if api in BotAccount.cookies:
cookies = BotAccount.cookies[api]
try:
return await get_url(api, status_code=200, headers=self.headers, fmt="json", request_private_ip=request_local)
return await get_url(api, status_code=200, headers=self.headers, fmt="json", request_private_ip=request_local,
cookies=cookies)
except Exception as e:
if api.find('moegirl.org.cn') != -1:
raise InvalidWikiError(self.locale.t("wiki.message.utils.wikilib.get_failed.moegirl"))