import re import time import os import random import json from base64 import b64decode from typing import Optional, List from datetime import datetime import logging import asyncio from io import BytesIO import bottom from PIL import Image, ImageDraw, ImageFont from PIL.JpegImagePlugin import JpegImageFile import aiohttp from settings import * TASK_QUEUE = asyncio.Queue() TASK_AUTHORS = {} FONT = ImageFont.truetype('font-bold.ttf', 22) bot = bottom.Client(host=IRC_HOST, port=IRC_PORT, ssl=IRC_TLS, loop=asyncio.get_event_loop()) USER_AGENT = 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/97.0.4692.99 Safari/537.36' HEADERS = { 'User-Agent': USER_AGENT, 'Accept': 'application/json', 'Accept-Language': 'en-US,en;q=0.5', 'Referer': 'https://www.craiyon.com/', 'Origin': 'https://www.craiyon.com', 'Connection': 'keep-alive', 'Sec-Fetch-Dest': 'empty', 'Sec-Fetch-Mode': 'cors', 'Sec-Fetch-Site': 'same-site', 'Pragma': 'no-cache', 'Cache-Control': 'no-cache', } class Task: def __init__(self, author: str, channel: str, term: str, dt: int): self.author = author self.channel = channel self.term = term self.dt = dt async def process_images(self, images: List[JpegImageFile]) -> JpegImageFile: if len(images) <= 0: raise Exception("Not enough images") if MODE == "monero": bg = Image.open('bg-pools.jpg') term_max = 60 term_x_start = 100 elif MODE == "2600nl": bg = Image.open('bg-2600nl.jpg') term_max = 52 term_x_start = 230 else: bg = Image.open('bg-wownero.jpg') term_max = 52 term_x_start = 230 d1 = ImageDraw.Draw(bg) term = self.term if len(term) >= term_max: term = self.term[:term_max] + ".." d1.text((term_x_start, 160), term, font=FONT, fill=(255, 255, 255)) x_start = 41 y_start = 232 for i, img in enumerate(images): x_offset = (i % 3) * (img.width + 20) y_offset = int(i / 3) * (img.width + 20) bg.paste(img, (x_start + x_offset, y_start + y_offset)) return bg async def post_image(self, image: BytesIO) -> str: url = 'https://uguu.se/upload.php' timeout = aiohttp.ClientTimeout(total=5) data = aiohttp.FormData() data.add_field('files[]', image, filename='lol.jpg', content_type='image/jpg') try: async with aiohttp.ClientSession( timeout=timeout, headers={"User-Agent": "w0w"}) as session: async with session.post(url, data=data) as resp: blob = await resp.json() return blob['files'][0]['url'] except Exception as ex: raise Exception("Could not upload file to uguu.se :((") async def get_images(self) -> Optional[List[JpegImageFile]]: url = 'https://backend.craiyon.com/generate' timeout_secs = 120 timeout = aiohttp.ClientTimeout(total=timeout_secs) images = [] try: async with aiohttp.ClientSession( timeout=timeout, headers=HEADERS) as session: async with session.post(url, json={'prompt': self.term}) as resp: if resp.status != 200: raise Exception("POST did not return status 200") blob = await resp.json() if 'images' not in blob: raise Exception("result did not return images") images = blob['images'] except asyncio.TimeoutError: raise Exception(f"image generation timeout exceeded ({timeout_secs} secs)") try: images = [b64decode(i) for i in images] return [Image.open(BytesIO(i)) for i in images] except Exception as ex: raise Exception("could not decode results") @bot.on('CLIENT_CONNECT') async def connect(**kwargs): bot.send('NICK', nick=IRC_NICK) bot.send('USER', user=IRC_NICK, realname=IRC_NICK) done, pending = await asyncio.wait( [bot.wait("RPL_ENDOFMOTD"), bot.wait("ERR_NOMOTD")], loop=bot.loop, return_when=asyncio.FIRST_COMPLETED ) for future in pending: future.cancel() for c in IRC_CHANNELS: bot.send('JOIN', channel=c) @bot.on('PING') def keepalive(message, **kwargs): bot.send('PONG', message=message) @bot.on('client_disconnect') def reconnect(**kwargs): logging.warning("Lost IRC server connection") time.sleep(3) bot.loop.create_task(bot.connect()) logging.warning("Reconnecting to IRC server") @bot.on('PRIVMSG') async def message_received(nick, target, message, **kwargs): if nick == IRC_NICK: return if target == IRC_NICK: target = nick if target not in IRC_CHANNELS: return msg = message msg = msg.replace("!dall", " !dall") msg = re.sub(r'!\W+', ' ', msg) now = datetime.now() now = time.mktime(now.timetuple()) await handle_msg(nick, msg, target, now) async def handle_msg(nick, msg, target, now: float): global TASK_AUTHORS, TASK_QUEUE def _err(target, err): return bot.send("PRIVMSG", target=target, message=err) if "re @" in msg: return if "!dall" not in msg: return msg = msg[msg.find("!dall"):] msg = msg[5:].strip() if len(msg) <= 8: return _err(target, f"longer query required.") TASK_AUTHORS.setdefault(nick, 0) if TASK_AUTHORS[nick] >= 3: err = f"{nick}: you already queued 3 thingies, patient!!11" bot.send("PRIVMSG", target=target, message=err) return TASK_AUTHORS[nick] += 1 task = Task(channel=target, author=nick, term=msg, dt=int(now)) await TASK_QUEUE.put(task) bot.send("PRIVMSG", target=target, message=f"{nick}: generating...") async def main(): global TASK_AUTHORS, TASK_QUEUE bot.loop.create_task(bot.connect()) def _lower_author_task_count(task_author): TASK_AUTHORS.setdefault(task_author, 1) TASK_AUTHORS[task_author] -= 1 def _err(_task, err): term = _task.term if len(term) >= 12: term = _task.term[:12] + ".." return bot.send("PRIVMSG", target=_task.channel, message=f"{_task.author}: Error \"{term}\"; {err}") # main loop while True: now = datetime.now() random.shuffle(TASK_QUEUE._queue) task = await TASK_QUEUE.get() try: images = await task.get_images() except Exception as ex: _err(task, err=f"Fetch: {ex}") _lower_author_task_count(task.author) continue try: image = await task.process_images(images) except Exception as ex: _err(task, err=f"Process: {ex}") _lower_author_task_count(task.author) continue image_bytes = BytesIO() image.save(image_bytes, 'JPEG') image.save(os.path.join('data', f"{task.dt}_{task.author}.jpg"), 'JPEG') image_bytes.seek(0) try: url = await task.post_image(image_bytes) except Exception as ex: _err(task, err=f"uguu upload: {ex}") _lower_author_task_count(task.author) continue _lower_author_task_count(task.author) completed_secs = int((datetime.now() - now).total_seconds()) bot.send("PRIVMSG", target=task.channel, message=f"{url} \"{task.term}\" (total: {completed_secs}s, {task.author}) ") await asyncio.sleep(10) loop = asyncio.get_event_loop() loop.run_until_complete(main())