import os import re import traceback import uuid from typing import Union from urllib.parse import urljoin import aiohttp import ujson as json from PIL import ImageFont from bs4 import BeautifulSoup, Comment from config import Config from core.logger import Logger from core.utils.http import download_to_cache web_render = Config('web_render') web_render_local = Config('web_render_local') elements = ['.notaninfobox', '.portable-infobox', '.infobox', '.tpl-infobox', '.infoboxtable', '.infotemplatebox', '.skin-infobox', '.arcaeabox', '.moe-infobox', '.rotable'] assets_path = os.path.abspath('./assets/') font = ImageFont.truetype(f'{assets_path}/SourceHanSansCN-Normal.ttf', 15) async def generate_screenshot_v2(page_link, section=None, allow_special_page=False, content_mode=False, use_local=True): elements_ = elements.copy() if not web_render_local: if not web_render: Logger.warn('[Webrender] Webrender is not configured.') return False use_local = False if section is None: if allow_special_page and content_mode: elements_.insert(0, '.mw-body-content') if allow_special_page and not content_mode: elements_.insert(0, '.diff') Logger.info('[Webrender] Generating element screenshot...') try: img = await download_to_cache((web_render_local if use_local else web_render) + 'element_screenshot', status_code=200, headers={'Content-Type': 'application/json'}, method="POST", post_data=json.dumps({ 'url': page_link, 'element': elements_}), attempt=1, timeout=30, request_private_ip=True ) except aiohttp.ClientConnectorError: if use_local: return await generate_screenshot_v2(page_link, section, allow_special_page, content_mode, use_local=False) else: return False except ValueError: Logger.info('[Webrender] Generation Failed.') return False else: Logger.info('[Webrender] Generating section screenshot...') try: img = await download_to_cache((web_render_local if use_local else web_render) + 'section_screenshot', status_code=200, headers={'Content-Type': 'application/json'}, method="POST", post_data=json.dumps({ 'url': page_link, 'section': section}), attempt=1, timeout=30, request_private_ip=True ) except aiohttp.ClientConnectorError: if use_local: return await generate_screenshot_v2(page_link, section, allow_special_page, content_mode, use_local=False) else: return False except ValueError: Logger.info('[Webrender] Generation Failed.') return False return img async def generate_screenshot_v1(link, page_link, headers, section=None, allow_special_page=False) -> Union[str, bool]: if not web_render_local: if not web_render: Logger.warn('[Webrender] Webrender is not configured.') return False try: Logger.info('Starting find infobox/section..') if link[-1] != '/': link += '/' try: async with aiohttp.ClientSession(headers=headers) as session: async with session.get(page_link, timeout=aiohttp.ClientTimeout(total=20)) as req: html = await req.read() except BaseException: traceback.print_exc() return False soup = BeautifulSoup(html, 'html.parser') pagename = uuid.uuid4() url = os.path.abspath(f'./cache/{pagename}.html') if os.path.exists(url): os.remove(url) Logger.info('Downloaded raw.') open_file = open(url, 'a', encoding='utf-8') timeless_fix = False def join_url(base, target): target = target.split(' ') targetlist = [] for x in target: if x.find('/') != -1: x = urljoin(base, x) targetlist.append(x) else: targetlist.append(x) target = ' '.join(targetlist) return target open_file.write('\n') for x in soup.find_all('html'): fl = [] for f in x.attrs: if isinstance(x.attrs[f], str): fl.append(f'{f}="{x.attrs[f]}"') elif isinstance(x.attrs[f], list): fl.append(f'{f}="{" ".join(x.attrs[f])}"') open_file.write(f'') open_file.write('
\n') for x in soup.find_all(rel='stylesheet'): if x.has_attr('href'): get_herf = x.get('href') if get_herf.find('timeless') != -1: timeless_fix = True x.attrs['href'] = re.sub(';', '&', urljoin(link, get_herf)) open_file.write(str(x)) for x in soup.find_all(): if x.has_attr('href'): x.attrs['href'] = re.sub(';', '&', urljoin(link, x.get('href'))) open_file.write('') for x in soup.find_all('style'): open_file.write(str(x)) if section is None: find_diff = None if allow_special_page: find_diff = soup.find('table', class_=re.compile('diff')) if find_diff is not None: Logger.info('Found diff...') for x in soup.find_all('body'): if x.has_attr('class'): open_file.write(f'') for x in soup.find_all('div'): if x.has_attr('id'): if x.get('id') in ['content', 'mw-content-text']: fl = [] for f in x.attrs: if isinstance(x.attrs[f], str): fl.append(f'{f}="{x.attrs[f]}"') elif isinstance(x.attrs[f], list): fl.append(f'{f}="{" ".join(x.attrs[f])}"') open_file.write(f'