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

headers support

Browse source
This commit is contained in:
yzhh 2021-03-21 13:00:17 +08:00
parent 76ad3135ef
commit 03ec82a83e
5 changed files with 240 additions and 151 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
modules/user/__init__.py
View file

@ -4,10 +4,12 @@ from graia.application import Group, Friend, MessageChain
from graia.application.message.elements.internal import Image, UploadMethods, Plain
from core.template import sendMessage
from modules.wiki.database import get_start_wiki, get_custom_interwiki
from modules.wiki.database import WikiDB
from .userlib import GetUser
database = WikiDB()
# 呜呜呜 想偷个懒都不行
async def main(kwargs: dict):
command = re.sub('^user ', '', kwargs['trigger_msg'])
@ -39,7 +41,7 @@ async def main(kwargs: dict):
table = 'custom_interwiki_group'
if Friend in kwargs:
table = 'custon_interwiki_self'
get_iw = get_custom_interwiki(table, id, match_interwiki.group(1))
get_iw = database.get_custom_interwiki(table, id, match_interwiki.group(1))
if get_iw:
metaurl = get_iw
username = match_interwiki.group(2)
@ -48,7 +50,7 @@ async def main(kwargs: dict):
table = 'start_wiki_link_group'
if Friend in kwargs:
table = 'start_wiki_link_self'
get_url = get_start_wiki(table, id)
get_url = database.get_start_wiki(table, id)
if get_url:
metaurl = get_url
username = command

82
modules/wiki/__init__.py
View file

@ -6,7 +6,7 @@ from graia.application.group import Group, Member
from graia.application.message.elements.internal import Image, UploadMethods
from graia.application.message.elements.internal import Plain
import modules.wiki.database as database
from modules.wiki.database import WikiDB
import modules.wiki.wikilib
from core.template import sendMessage, check_permission, wait_confirm, revokeMessage
from database import warn_someone, check_enable_modules_self, check_enable_modules
@ -14,13 +14,42 @@ from modules.wiki.helper import check_wiki_available
from .getinfobox import get_infobox_pic
database = WikiDB()
async def wiki_loader(kwargs: dict):
kwargs['trigger_msg'] = cmd = re.sub(r'^wiki ', '', kwargs['trigger_msg'])
cmd = cmd.split(' ')
if isinstance(cmd, list):
if len(cmd) > 1:
if cmd[0] == 'set':
kwargs['trigger_msg'] = cmd[1]
await set_start_wiki(kwargs)
if cmd[0] == 'iw':
kwargs['trigger_msg'] = ' '.join(cmd[1:])
await interwiki(kwargs)
if cmd[0] == 'headers':
kwargs['trigger_msg'] = ' '.join(cmd[1:])
await set_headers(kwargs)
else:
await wiki_wrapper(kwargs)
else:
await wiki_wrapper(kwargs)
async def wiki_wrapper(kwargs: dict):
command = kwargs['trigger_msg']
command = re.sub(r'^wiki ', '', command)
if Group in kwargs:
start_table = 'start_wiki_link_group'
headtable = 'request_headers_group'
headtarget = kwargs[Group].id
if Friend in kwargs:
start_table = 'start_wiki_link_self'
headtable = 'request_headers_self'
headtarget = kwargs[Friend].id
headers = database.config_headers('get', headtable, headtarget)
prompt = False
get_link = database.get_start_wiki(start_table, kwargs[Group].id)
if not get_link:
@ -108,7 +137,7 @@ async def wiki_loader(kwargs: dict):
iw = matchinterwiki.group(1)
get_link = get_custom_iw
command = re.sub(matchinterwiki.group(1) + ':', '', command)
msg = await wikilib.wikilib().main(get_link, command, interwiki=iw)
msg = await wikilib.wikilib().main(get_link, command, interwiki=iw, headers=headers)
if msg['status'] == 'done':
msgchain = MessageChain.create(
[Plain((prompt + '\n' if prompt else '') + (msg['url'] + '\n' if 'url' in msg else '') + msg['text'])])
@ -130,12 +159,8 @@ async def wiki_loader(kwargs: dict):
'wiki_infobox')
print(check_options)
if check_options:
pic = await get_infobox_pic(get_link, msg['url'])
if Group in kwargs:
mth = UploadMethods.Group
elif Friend in kwargs:
mth = UploadMethods.Friend
imgchain = MessageChain.create([Image.fromLocalFile(pic, method=mth)])
pic = await get_infobox_pic(get_link, msg['url'], headers)
imgchain = MessageChain.create([Image.fromLocalFile(pic)])
await sendMessage(kwargs, imgchain)
elif msg['status'] == 'wait':
@ -154,10 +179,10 @@ async def wiki_loader(kwargs: dict):
async def set_start_wiki(kwargs: dict):
command = kwargs['trigger_msg']
command = re.sub(r'^wiki_start_site ', '', command)
if Group in kwargs:
if check_permission(kwargs):
command = kwargs['trigger_msg']
command = re.sub(r'^wiki_start_site ', '', command)
check = await check_wiki_available(command)
if check:
result = database.add_start_wiki('start_wiki_link_group', kwargs[Group].id, check[0])
@ -169,8 +194,6 @@ async def set_start_wiki(kwargs: dict):
result = '你没有使用该命令的权限。'
await sendMessage(kwargs, MessageChain.create([Plain(result)]))
if Friend in kwargs:
command = kwargs['trigger_msg']
command = re.sub(r'^wiki_start_site ', '', command)
check = await check_wiki_available(command)
if check:
result = database.add_start_wiki('start_wiki_link_self', kwargs[Friend].id, check[0])
@ -184,6 +207,7 @@ async def interwiki(kwargs: dict):
command = kwargs['trigger_msg']
command = re.sub(r'^interwiki ', '', command)
command = command.split(' ')
print(command)
if Group in kwargs:
check = check_permission(kwargs)
if not check:
@ -196,7 +220,7 @@ async def interwiki(kwargs: dict):
table = 'custom_interwiki_self'
target = kwargs[Friend].id
if command[0] == 'add':
command = re.sub(r'^interwiki add ', '', kwargs['trigger_msg'])
command = ' '.join(command[1:])
command = re.sub(' ', '>', command)
iw = command.split('>')
try:
@ -207,7 +231,7 @@ async def interwiki(kwargs: dict):
if check:
result = database.config_custom_interwiki('add', table, target, iw[0],
check[0])
await sendMessage(kwargs, MessageChain.create([Plain(result + check[1])]))
await sendMessage(kwargs, MessageChain.create([Plain(result + f'{iw[0]} > {check[1]}')]))
else:
result = '错误：此Wiki不是一个有效的MediaWiki/尝试建立连接超时。'
link = re.match(r'^(https?://).*', iw[1])
@ -231,6 +255,29 @@ async def interwiki(kwargs: dict):
await sendMessage(kwargs, '命令不合法，参数应为add/del/list。')
async def set_headers(kwargs: dict):
command = kwargs['trigger_msg']
command = command.split(' ')
if Group in kwargs:
check = check_permission(kwargs)
if not check:
result = '你没有使用该命令的权限。'
await sendMessage(kwargs, MessageChain.create([Plain(result)]))
return
table = 'request_headers_group'
id = kwargs[Group].id
if Friend in kwargs:
table = 'request_headers_self'
id = kwargs[Friend].id
do = command[0]
if do == 'show':
headers = database.config_headers(do, table, id)
msg = f'当前设置了以下标头：\n{headers}\n如需自定义，请使用~wiki headers <set> <headers>，不同标头之间使用换行隔开。'
else:
msg = database.config_headers(do, table, id, ' '.join(command[1:]))
await sendMessage(kwargs, msg)
async def regex_wiki(kwargs: dict):
display = kwargs[MessageChain].asDisplay()
@ -259,10 +306,13 @@ async def regex_wiki(kwargs: dict):
table = 'start_wiki_link_group'
target = kwargs[Group].id
mth = UploadMethods.Group
headtable = 'request_headers_group'
if Friend in kwargs:
table = 'start_wiki_link_self'
target = kwargs[Friend].id
mth = UploadMethods.Friend
headtable = 'request_headers_self'
headers = database.config_headers('get', headtable, target)
for find in find_dict:
if find_dict[find] == 'template':
template = True
@ -315,7 +365,7 @@ async def regex_wiki(kwargs: dict):
get_link = f'https://{matchinterwiki.group(1)}.fandom.com/api.php'
find = matchinterwiki.group(2)
iw = matchinterwiki.group(1)
msg = await modules.wiki.wikilib.wikilib().main(get_link, find, interwiki=iw, template=template)
msg = await modules.wiki.wikilib.wikilib().main(get_link, find, interwiki=iw, template=template, headers=headers)
status = msg['status']
text = (prompt + '\n' if prompt else '') + msg['text']
if status == 'wait':

269
modules/wiki/database.py
View file

@ -1,134 +1,171 @@
import os
import re
import sqlite3
import base64
import traceback
dbpath = os.path.abspath('./modules/wiki/save.db')
def initialize():
a = open(dbpath, 'w')
a.close()
conn = sqlite3.connect(dbpath)
c = conn.cursor()
c.execute('''CREATE TABLE start_wiki_link_group
(ID INT PRIMARY KEY NOT NULL,
LINK TEXT);''')
c.execute('''CREATE TABLE custom_interwiki_group
(ID INT PRIMARY KEY NOT NULL,
INTERWIKIS TEXT);''')
c.execute('''CREATE TABLE start_wiki_link_self
(ID INT PRIMARY KEY NOT NULL,
LINK TEXT);''')
c.execute('''CREATE TABLE custom_interwiki_self
(ID INT PRIMARY KEY NOT NULL,
INTERWIKIS TEXT);''')
c.close()
class WikiDB:
def __init__(self):
if not os.path.exists(dbpath):
self.initialize()
self.conn = sqlite3.connect(dbpath)
self.c = self.conn.cursor()
def add_start_wiki(table, id, value):
if not os.path.exists(dbpath):
initialize()
conn = sqlite3.connect(dbpath)
c = conn.cursor()
a = c.execute(f"SELECT * FROM {table} WHERE ID={id}").fetchone()
if a:
c.execute(f"UPDATE {table} SET LINK='{value}' WHERE ID='{id}'")
conn.commit()
return '成功设置起始Wiki：'
else:
c.execute(f"INSERT INTO {table} (ID, Link) VALUES (?, ?)", (id, value))
conn.commit()
return '成功设置起始Wiki：'
def initialize(self):
a = open(dbpath, 'w')
a.close()
self.conn = sqlite3.connect(dbpath)
self.c = self.conn.cursor()
self.c.execute('''CREATE TABLE start_wiki_link_group
(ID INT PRIMARY KEY NOT NULL,
LINK TEXT);''')
self.c.execute('''CREATE TABLE custom_interwiki_group
(ID INT PRIMARY KEY NOT NULL,
INTERWIKIS TEXT);''')
self.c.execute('''CREATE TABLE start_wiki_link_self
(ID INT PRIMARY KEY NOT NULL,
LINK TEXT);''')
self.c.execute('''CREATE TABLE custom_interwiki_self
(ID INT PRIMARY KEY NOT NULL,
INTERWIKIS TEXT);''')
self.c.execute('''CREATE TABLE wiki_info
(LINK TEXT PRIMARY KEY NOT NULL,
SITEINFO TEXT,
TS TIMESTAMP DEFAULT CURRENT_TIMESTAMP);''')
self.c.execute('''CREATE TABLE request_headers_group
(ID TEXT PRIMARY KEY NOT NULL,
HEADERS TEXT);''')
self.c.execute('''CREATE TABLE request_headers_self
(ID TEXT PRIMARY KEY NOT NULL,
HEADERS TEXT);''')
self.c.close()
def get_start_wiki(table, id):
if not os.path.exists(dbpath):
initialize()
conn = sqlite3.connect(dbpath)
c = conn.cursor()
a = c.execute(f"SELECT * FROM {table} WHERE ID={id}").fetchone()
if a:
return a[1]
else:
return False
def config_custom_interwiki(do, table, id, iw, link=None):
if not os.path.exists(dbpath):
initialize()
conn = sqlite3.connect(dbpath)
c = conn.cursor()
a = c.execute(f"SELECT * FROM {table} WHERE ID={id}").fetchone()
if do == 'add':
def add_start_wiki(self, table, id, value):
a = self.c.execute(f"SELECT * FROM {table} WHERE ID={id}").fetchone()
if a:
split_iws = a[1].split('|')
iwlist = []
for iws in split_iws:
split_iw = iws.split('>')
iwlist.append(split_iw[0])
if iw in iwlist:
self.c.execute(f"UPDATE {table} SET LINK='{value}' WHERE ID='{id}'")
self.conn.commit()
return '成功设置起始Wiki：'
else:
self.c.execute(f"INSERT INTO {table} (ID, Link) VALUES (?, ?)", (id, value))
self.conn.commit()
return '成功设置起始Wiki：'
def get_start_wiki(self, table, id):
a = self.c.execute(f"SELECT * FROM {table} WHERE ID={id}").fetchone()
if a:
return a[1]
else:
return False
def config_custom_interwiki(self, do, table, id, iw, link=None):
a = self.c.execute(f"SELECT * FROM {table} WHERE ID={id}").fetchone()
if do == 'add':
if a:
split_iws = a[1].split('|')
iwlist = []
for iws in split_iws:
if iws.find(iw + '>') != -1:
split_iws.remove(iws)
split_iws.append(f'{iw}>{link}')
c.execute(
f"UPDATE {table} SET INTERWIKIS='{'|'.join(split_iws)}' WHERE ID='{id}'")
conn.commit()
return '成功：更新自定义Interwiki：'
split_iw = iws.split('>')
iwlist.append(split_iw[0])
if iw in iwlist:
for iws in split_iws:
if iws.find(iw + '>') != -1:
split_iws.remove(iws)
split_iws.append(f'{iw}>{link}')
self.c.execute(
f"UPDATE {table} SET INTERWIKIS='{'|'.join(split_iws)}' WHERE ID='{id}'")
self.conn.commit()
return '成功：更新自定义Interwiki：'
else:
split_iws.append(f'{iw}>{link}')
self.c.execute(
f"UPDATE {table} SET INTERWIKIS='{'|'.join(split_iws)}' WHERE ID='{id}'")
self.conn.commit()
return '成功：添加自定义Interwiki：'
else:
split_iws.append(f'{iw}>{link}')
c.execute(
f"UPDATE {table} SET INTERWIKIS='{'|'.join(split_iws)}' WHERE ID='{id}'")
conn.commit()
self.c.execute(f"INSERT INTO {table} (ID, INTERWIKIS) VALUES (?, ?)", (id, f'{iw}>{link}'))
self.conn.commit()
return '成功：添加自定义Interwiki：'
else:
c.execute(f"INSERT INTO {table} (ID, INTERWIKIS) VALUES (?, ?)", (id, f'{iw}>{link}'))
conn.commit()
return '成功：添加自定义Interwiki：'
elif do == 'del':
if a:
split_iws = a[1].split('|')
iwlist = []
for iws in split_iws:
split_iw = iws.split('>')
iwlist.append(split_iw[0])
if iw in iwlist:
elif do == 'del':
if a:
split_iws = a[1].split('|')
iwlist = []
for iws in split_iws:
if iws.find(iw + '>') != -1:
split_iws.remove(iws)
c.execute(
f"UPDATE {table} SET INTERWIKIS='{'|'.join(split_iws)}' WHERE ID='{id}'")
conn.commit()
return '成功：删除自定义Interwiki：'
split_iw = iws.split('>')
iwlist.append(split_iw[0])
if iw in iwlist:
for iws in split_iws:
if iws.find(iw + '>') != -1:
split_iws.remove(iws)
self.c.execute(
f"UPDATE {table} SET INTERWIKIS='{'|'.join(split_iws)}' WHERE ID='{id}'")
self.conn.commit()
return '成功：删除自定义Interwiki：'
else:
return '失败：添加过此Interwiki：'
else:
return '失败：添加过此Interwiki：'
return '失败：未添加过任何Interwiki。'
def get_custom_interwiki(self, table, id, iw):
a = self.c.execute(f"SELECT * FROM {table} WHERE ID={id}").fetchone()
if a:
interwikis = a[1].split('|')
for iws in interwikis:
if iws.find(iw + '>') != -1:
iws = iws.split('>')
return iws[1]
else:
return '失败：未添加过任何Interwiki。'
return False
def get_custom_interwiki_list(self, table, id):
a = self.c.execute(f"SELECT * FROM {table} WHERE ID={id}").fetchone()
if a:
interwikis = a[1].split('|')
return '\n'.join(interwikis)
else:
return False
def config_headers(self, do, table, id, header=None):
try:
a = self.c.execute(f"SELECT * FROM {table} WHERE ID={id}").fetchone()
if do == 'set':
header = base64.encodebytes(header.encode('utf-8'))
if a:
self.c.execute(f"UPDATE {table} SET HEADERS='{header}' WHERE ID='{id}'")
else:
self.c.execute(f"INSERT INTO {table} (ID, HEADERS) VALUES (?, ?)", (id, header))
self.conn.commit()
return '成功更新请求所使用的Headers。'
elif do == 'reset':
if a:
self.c.execute(f"DELETE FROM {table} WHERE ID='{id}'")
return '成功重置请求所使用的Headers。'
else:
return '当前未自定义请求所使用的Headers。'
elif do == 'get':
headers = {}
if a:
headerd = str(base64.decodebytes(a[1]).decode('utf-8'))
headersp = headerd.split('\n')
for x in headersp:
x = x.split(':')
headers[x[0]] = re.sub(r'^ ', '', x[1])
else:
headers['accept-language'] = 'zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6'
print(header)
return headers
elif do == 'show':
if a:
header = str(base64.decodebytes(a[1]).decode('utf-8'))
else:
header = '（默认）accept-language: zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6'
return header
except Exception as e:
traceback.print_exc()
return '发生错误' + str(e)
def get_custom_interwiki(table, id, iw):
if not os.path.exists(dbpath):
initialize()
conn = sqlite3.connect(dbpath)
c = conn.cursor()
a = c.execute(f"SELECT * FROM {table} WHERE ID={id}").fetchone()
if a:
interwikis = a[1].split('|')
for iws in interwikis:
if iws.find(iw + '>') != -1:
iws = iws.split('>')
return iws[1]
else:
return False
def get_custom_interwiki_list(table, id):
if not os.path.exists(dbpath):
initialize()
conn = sqlite3.connect(dbpath)
c = conn.cursor()
a = c.execute(f"SELECT * FROM {table} WHERE ID={id}").fetchone()
if a:
interwikis = a[1].split('|')
return '\n'.join(interwikis)
else:
return False

5
modules/wiki/getinfobox.py
View file

@ -12,7 +12,6 @@ from selenium import webdriver
from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
from config import config
from modules.wiki.helper import get_url
config_path = os.path.abspath('./config/config.cfg')
@ -23,7 +22,7 @@ except:
infobox_render = None
async def get_infobox_pic(link, pagelink):
async def get_infobox_pic(link, pagelink, headers):
try:
print('hello')
wlink = re.sub(r'api.php', '', link)
@ -31,7 +30,7 @@ async def get_infobox_pic(link, pagelink):
print(link)
print(pagelink)
try:
async with aiohttp.ClientSession() as session:
async with aiohttp.ClientSession(headers=headers) as session:
async with session.get(pagelink, timeout=aiohttp.ClientTimeout(total=20)) as req:
html = await req.read()
except:

27
modules/wiki/wikilib.py
View file

@ -9,8 +9,8 @@ from .helper import check_wiki_available
class wikilib:
async def get_data(self, url: str, fmt: str):
async with aiohttp.ClientSession() as session:
async def get_data(self, url: str, fmt: str, headers=None):
async with aiohttp.ClientSession(headers=headers) as session:
try:
async with session.get(url, timeout=aiohttp.ClientTimeout(total=20)) as req:
if hasattr(req, fmt):
@ -72,7 +72,6 @@ class wikilib:
d[j['query']['namespaces'][x]['*']] = j['query']['namespaces'][x]['canonical']
except:
traceback.print_exc()
pass
return d
async def get_article_path(self, url):
@ -111,13 +110,13 @@ class wikilib:
try:
try:
searchurl = self.wikilink + '?action=query&generator=search&gsrsearch=' + self.pagename + '&gsrsort=just_match&gsrenablerewrites&prop=info&gsrlimit=1&format=json'
getsecjson = await self.get_data(searchurl, "json")
getsecjson = await self.get_data(searchurl, "json", self.headers)
secpageid = self.parsepageid(getsecjson)
sectitle = getsecjson['query']['pages'][secpageid]['title']
except:
traceback.print_exc()
searchurl = self.wikilink + '?action=query&list=search&srsearch=' + self.pagename + '&srwhat=text&srlimit=1&srenablerewrites=&format=json'
getsecjson = await self.get_data(searchurl, "json")
getsecjson = await self.get_data(searchurl, "json", self.headers)
sectitle = getsecjson['query']['search'][0]['title']
if self.interwiki == '':
target = ''
@ -158,7 +157,7 @@ class wikilib:
async def getdesc(self):
try:
descurl = self.wikilink + '?action=query&prop=info|pageprops|extracts&ppprop=description|displaytitle|disambiguation|infoboxes&explaintext=true&exsectionformat=plain&exsentences=1&format=json&titles=' + self.querytextname
loadtext = await self.get_data(descurl, "json")
loadtext = await self.get_data(descurl, "json", self.headers)
pageid = self.parsepageid(loadtext)
desc = loadtext['query']['pages'][pageid]['extract']
except Exception:
@ -169,7 +168,7 @@ class wikilib:
async def getfirstline(self):
try:
descurl = self.wikilink + f'?action=parse&page={self.querytextname}&prop=wikitext&section=0&format=json'
loaddesc = await self.get_data(descurl, 'json')
loaddesc = await self.get_data(descurl, 'json', self.headers)
descraw = loaddesc['parse']['wikitext']['*']
try:
cutdesc = re.findall(r'(.*(?:!|\?|\.|;|！|？|。|；))', descraw, re.S | re.M)
@ -184,7 +183,7 @@ class wikilib:
async def getalltext(self):
try:
descurl = self.wikilink + f'?action=parse&page={self.querytextname}&prop=wikitext&format=json'
loaddesc = await self.get_data(descurl, 'json')
loaddesc = await self.get_data(descurl, 'json', self.headers)
desc = loaddesc['parse']['wikitext']['*']
except Exception:
traceback.print_exc()
@ -256,14 +255,14 @@ class wikilib:
finpgname = geturlpagename + urllib.parse.quote(section.group(1).encode('UTF-8'))
fullurl = self.psepgraw['fullurl'] + urllib.parse.quote(section.group(1).encode('UTF-8'))
except Exception:
pass
traceback.print_exc()
try:
pgtag = re.match(r'.*(\?.*)', self.pagename)
if pgtag:
finpgname = geturlpagename + pgtag.group(1)
fullurl = fullurl + pgtag.group(1)
except Exception:
pass
traceback.print_exc()
finpgname = urllib.parse.unquote(finpgname)
finpgname = re.sub('_', ' ', finpgname)
if finpgname == self.orginpagename:
@ -291,14 +290,15 @@ class wikilib:
getimg = await self.get_image(self.pagename)
if getimg:
msgs['net_image'] = getimg
if await self.danger_text_check(result):
print(result)
if result != '' and await self.danger_text_check(result):
return {'status': 'done', 'text': 'https://wdf.ink/6OUp'}
return msgs
except Exception as e:
traceback.print_exc()
return {'status': 'done', 'text': '发生错误：' + str(e)}
async def main(self, wikilink, pagename, interwiki=None, igmessage=False, template=False, tryiw=0):
async def main(self, wikilink, pagename, interwiki=None, igmessage=False, template=False, headers=None, tryiw=0):
print(wikilink)
print(pagename)
print(interwiki)
@ -321,6 +321,7 @@ class wikilib:
self.igmessage = igmessage
self.template = template
self.templateprompt = None
self.headers = headers
try:
matchinterwiki = re.match(r'(.*?):(.*)', self.pagename)
if matchinterwiki:
@ -335,7 +336,7 @@ class wikilib:
((
interwiki + ':') if interwiki is not None else '') + matchinterwiki.group(
1),
self.igmessage, self.template, tryiw + 1)
self.igmessage, self.template, headers, tryiw + 1)
else:
return {'status': 'done',
'text': f'发生错误：指向的interwiki不是一个有效的MediaWiki。{interwiki_link}{matchinterwiki.group(2)}'}