rollbot4/commands/commands/querying.py

from logging import Logger
import random
import asyncio

from aiohttp import MultipartWriter
from bs4 import BeautifulSoup

from rollbot import as_command, RollbotFailure, Attachment
from rollbot.injection import Request, Args, Arg, Config


@as_command
async def inspire(req: Request, logger: Logger):
    """
    Using the command !inspire will request an inspiration image from http://inspirobot.me/
    """
    try:
        async with req.get("http://inspirobot.me/api?generate=true") as res:
            return await res.text()
    except:
        logger.exception("Failed !inspire")
        RollbotFailure.SERVICE_DOWN.raise_exc(detail="Could not reach inspirobot")


@as_command
async def shield(req: Request, blazon: Args, logger: Logger):
    boundary = f"----RollbotBoundary{''.join(str(random.randint(0, 9)) for _ in range(10))}"
    headers = {
        "Content-Type": f"multipart/form-data; boundary={boundary}",
        "Cache-Control": "no-cache",
    }
    with MultipartWriter(boundary=boundary) as form:
        form.append(blazon).set_content_disposition("form-data", name="blazon")
        form.append("1").set_content_disposition("form-data", name="asfile")
        form.append("wappenwiki").set_content_disposition("form-data", name="palette")
        form.append("flat").set_content_disposition("form-data", name="effect")
        form.append("500").set_content_disposition("form-data", name="size")
    try:
        async with req.post(
            "https://drawshield.net/include/drawshield.php",
            headers=headers,
            data=form,
        ) as res:
            res.raise_for_status()
            return Attachment("image", await res.read())
    except:
        logger.exception("Failed !shield")
        RollbotFailure.SERVICE_DOWN.raise_exc(detail="Could not reach DrawShield")


@as_command
async def scp(
    number: Arg(0, convert=int, fail_msg="Could not parse argument {} into integer"), req: Request
):
    page_url = f"http://www.scp-wiki.net/scp-{number:03d}"
    series_url = "http://www.scp-wiki.net/scp-series"
    series = (number // 1000) + 1
    if series != 1:
        series_url += "-" + str(series)

    async def get_obj_class():
        async with req.get(page_url) as page:
            page_html = BeautifulSoup(await page.text(), "html.parser")

        for p in page_html.find_all("p"):
            if "Object Class" in p.text:
                return p.text

        return "Error retrieving object class!"

    query = f"SCP-{number:03d}"

    async def get_title():
        async with req.get(series_url) as series:
            series_html = BeautifulSoup(await series.text(), "html.parser")

        for li in series_html.find_all("li"):
            if query in li.text:
                return li.text.split("-", 2)[-1].strip()

        return "Error retrieving title!"

    obj_class, title = await asyncio.gather(get_obj_class(), get_title())
    return f"Item # {number}\n{obj_class}\n{title}\n{page_url}"


@as_command
async def riddle(args: Args, req: Request, logger: Logger, sleep: Config("riddle.sleep")):
    if args.lower() != "me piss":
        RollbotFailure.INVALID_ARGUMENTS.raise_exc()

    rnum = None
    riddle = None
    for n in range(20):
        rnum = random.randint(2, 10050)
        async with req.get(f"https://www.riddles.com/{rnum}") as res:
            riddle_page = BeautifulSoup(await res.text(), "html.parser")

        riddle_panel = riddle_page.find("div", class_="panel-body")
        if riddle_panel is not None:
            parts = [p.text for p in riddle_panel.find_all("p")]
            if (
                len(parts) == 2
                and not parts[0].startswith("Riddle Status: User Rejected")
                and not parts[0].startswith("Riddle Status: User Moderated")
            ):
                riddle = parts[0]
                answer = parts[1]
                break
        logger.info(f"Failed to parse riddle #{rnum}, attempt #{n}")

    else:
        RollbotFailure.SERVICE_DOWN.raise_exc(detail="Failed to find a riddle fast enough")

    yield "\n\n".join(
        (
            f"Riddle #{rnum}",
            riddle,
            f"I'll post the response in about {sleep} seconds!",
        )
    )

    # TODO need to confirm this works without blocking in a proper task deploying system
    await asyncio.sleep(sleep)

    yield "\n\n".join(("Here's the riddle answer from before!", riddle, answer))


@as_command
async def selfie(req: Request, logger: Logger):
    try:
        async with req.get(
            "https://cdn.star.nesdis.noaa.gov/GOES16/ABI/SECTOR/ne/GEOCOLOR/latest.jpg"
        ) as res:
            res.raise_for_status()
            return Attachment("image", await res.read())
    except:
        logger.exception("Failed to call GOES16")
        RollbotFailure.SERVICE_DOWN.raise_exc(detail="Failed to call GOES16")


@as_command
async def cat(req: Request, logger: Logger):
    """
    The !cat command grabs a cat from https://thiscatdoesnotexist.com/
    """
    try:
        async with req.get(
            "https://thiscatdoesnotexist.com/", headers={"User-Agent": "Rollbot"}
        ) as res:
            res.raise_for_status()
            return Attachment("image", await res.read())
    except:
        logger.exception("Failed to call cat generator")
        RollbotFailure.SERVICE_DOWN.raise_exc(detail="Failed to call cat generator")


@as_command
async def horse(req: Request, logger: Logger):
    """
    The !horse command grabs a horse from https://thishorsedoesnotexist.com/
    """
    try:
        async with req.get(
            "https://thishorsedoesnotexist.com/", headers={"User-Agent": "Rollbot"}
        ) as res:
            res.raise_for_status()
            return Attachment("image", await res.read())
    except:
        logger.exception("Failed to call horse generator")
        RollbotFailure.SERVICE_DOWN.raise_exc(detail="Failed to call horse generator")


@as_command
async def npc(req: Request, logger: Logger):
    """
    The !horse command grabs a person from https://thispersondoesnotexist.com/
    """
    try:
        async with req.get(
            "https://thispersondoesnotexist.com/", headers={"User-Agent": "Rollbot"}
        ) as res:
            res.raise_for_status()
            return Attachment("image", await res.read())
    except:
        logger.exception("Failed to call person generator")
        RollbotFailure.SERVICE_DOWN.raise_exc(detail="Failed to call person generator")


art_number = 0


@as_command
async def art(req: Request, logger: Logger):
    """
    The !art command uses the 9gans gallery at https://9gans.com/ which generates 9 images every hour.
    This command will cycle through those 9 images, so if you fire it 10 times in quick succession, the tenth
    piece of art might be the same as the first.
    """
    global art_number
    art_number += 1
    art_number %= 9
    try:
        async with req.get(
            f"https://storage.googleapis.com/9gans/mini/{art_number + 1}.jpg"
        ) as res:
            res.raise_for_status()
            return Attachment("image", await res.read())
    except:
        logger.exception("Failed to call art generator")
        RollbotFailure.SERVICE_DOWN.raise_exc(detail="Failed to call art generator")


@as_command
async def imagine(args: Args, api_key: Config("text2img.api_key"), req: Request, logger: Logger):
    """
    The !imagine command uses the text2img API at https://deepai.org/machine-learning-model/text2img
    """

    args = args.strip()

    if len(args) == 0:
        RollbotFailure.INVALID_ARGUMENTS.raise_exc(
            detail="The !imagine command needs text to imagine!"
        )

    try:
        async with req.post(
            "https://api.deepai.org/api/text2img", data={"text": args}, headers={"api-key": api_key}
        ) as res:
            res.raise_for_status()
            print(await res.text())
            js = await res.json()
            return js.get("output_url", None)
    except:
        logger.exception("Failed to call text2img service")
        RollbotFailure.SERVICE_DOWN.raise_exc(detail="Failed to call text2img service")