diff --git a/core/utils/http.py b/core/utils/http.py index d171af5a..e8d9cc76 100644 --- a/core/utils/http.py +++ b/core/utils/http.py @@ -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), diff --git a/modules/wiki/audit.py b/modules/wiki/audit.py index 1acf7bef..045dda2b 100644 --- a/modules/wiki/audit.py +++ b/modules/wiki/audit.py @@ -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 ') +async def _(msg: Bot.MessageSession): + api_link = msg.parsed_msg[''] + account = msg.parsed_msg[''] + password = msg.parsed_msg[''] + 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 ') +async def _(msg: Bot.MessageSession): + api_link = msg.parsed_msg[''] + 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() diff --git a/modules/wiki/utils/bot.py b/modules/wiki/utils/bot.py new file mode 100644 index 00000000..40976a87 --- /dev/null +++ b/modules/wiki/utils/bot.py @@ -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}') diff --git a/modules/wiki/utils/dbutils.py b/modules/wiki/utils/dbutils.py index 556f4996..ca5bec68 100644 --- a/modules/wiki/utils/dbutils.py +++ b/modules/wiki/utils/dbutils.py @@ -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() diff --git a/modules/wiki/utils/orm.py b/modules/wiki/utils/orm.py index ca13fabc..cb2c1bf8 100644 --- a/modules/wiki/utils/orm.py +++ b/modules/wiki/utils/orm.py @@ -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() diff --git a/modules/wiki/utils/wikilib.py b/modules/wiki/utils/wikilib.py index 84f5b46b..b463db99 100644 --- a/modules/wiki/utils/wikilib.py +++ b/modules/wiki/utils/wikilib.py @@ -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"))