update
This commit is contained in:
parent
7dac27ccff
commit
1660118cc6
6 changed files with 130 additions and 5 deletions
|
@ -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),
|
||||
|
|
|
@ -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
43
modules/wiki/utils/bot.py
Normal 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}')
|
|
@ -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()
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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"))
|
||||
|
|
Reference in a new issue