rollbot3/src/plugins/querying.py

import random
import threading
import time

import requests
from requests_html import HTMLSession

from command_system import RollbotResponse, RollbotFailure, as_plugin, pop_arg
from config import get_config, get_secret
from util import upload_image

@as_plugin
def inspire(message):
    argstr = message.raw_args

    if argstr is not None and "explain" in argstr:
        return "Using the command !inspire will request an inspiration image from http://inspirobot.me/"

    try:
        url = requests.get("http://inspirobot.me/api?generate=true").text
        return RollbotResponse(message, img=url)
    except ConnectionError as e:
        return RollbotResponse(
            message,
            failure=RollbotFailure.SERVICE_DOWN,
            debugging={
                "explain": "Could not reach inspirobot.",
                "exception": e
            }
        )


@as_plugin
def shield(message):
    data = """\
------RollbotBoundary\r\n\
Content-Disposition: form-data; name=\"blazon\"\r\n\
\r\n\
%s\r\n\
------RollbotBoundary\r\n\
Content-Disposition: form-data; name=\"asfile\"\r\n\
\r\n\
1\r\n\
------RollbotBoundary\r\n\
Content-Disposition: form-data; name=\"palette\"\r\n\
\r\n\
wappenwiki\r\n\
------RollbotBoundary\r\n\
Content-Disposition: form-data; name=\"effect\"\r\n\
\r\n\
flat\r\n\
------RollbotBoundary
""" % message.raw_args
    # TODO - this should probably have a random number as part of the boundary

    try:
        r = requests.post(
            "https://drawshield.net/include/drawshield.php",
            headers={
                'Content-Type': "multipart/form-data; boundary=----RollbotBoundary",
                'Cache-Control': "no-cache",
            },
            data=data
        )
    except ConnectionError as e:
        return RollbotResponse(
            message,
            failure=RollbotFailure.SERVICE_DOWN,
            debugging={
                "explain": "Could not reach DrawShield.",
                "exception": e,
            }
        )

    success, result = upload_image(get_secret("imgur_client_id"), r.content)
    if success:
        return RollbotResponse(message, img=result)
    else:
        return RollbotResponse(
            message,
            failure=RollbotFailure.SERVICE_DOWN,
            debugging=result
        )


@as_plugin
def scp(message):
    number, _ = pop_arg(message.raw_args)
    try:
        number = int(number.strip())
    except ValueError as e:
        return RollbotResponse(
            message,
            failure=RollbotFailure.INVALID_ARGUMENTS,
            debugging={
                "explain": "Could not parse argument %s into integer." % number,
                "exception": e,
            }
        )

    page_url = "http://www.scp-wiki.net/scp-%03d" % number
    series_url = "http://www.scp-wiki.net/scp-series"
    series = (number // 1000) + 1
    if series != 1:
        series_url += "-" + str(series)

    sesh = HTMLSession()

    r = sesh.get(page_url)
    r2 = sesh.get(series_url)

    ps = r.html.find("p")
    lis = r2.html.find("li")

    object_class = "Error retrieving object class!"
    for p in ps:
        if "Object Class" in p.text:
            object_class = p.text
            break

    query = "SCP-%03d" % number
    flavortext = "Error retrieving title!"
    for li in lis:
        if query in li.text:
            flavortext = li.text.split("-", 2)[-1].strip()
            break

    response = "Item # %d\n" % number
    response += object_class + "\n"
    response += flavortext + "\n"
    response += page_url

    return response


def get_riddle(sesh, logger):
    n = random.randint(2, 7545)
    r = sesh.get(f"https://www.riddles.com/{n:d}")
    divs = r.html.find("div")
    panel_bodies = [d.text for d in divs if "class" in d.attrs and "panel-body" in d.attrs["class"]]

    if len(panel_bodies) <= 1:
        logger.info(f"Failed on {n:d} - invalid page")
        return n, None

    riddle = panel_bodies[0]

    if riddle.startswith("Riddle Status: User Rejected") or riddle.startswith("Riddle Status: User Moderated"):
        logger.info(f"Failed on {n:d} - rejected")
        return n, None

    parts = riddle.split("\n")
    logger.info(f"Returning riddle {n:d}")
    return n, (parts[1], parts[3])


def get_answer_cb(bot, message, riddle):
    gid = message.group_id
    txt = "\n\n".join(("Here's the riddle answer from before!", *riddle))
    def cb():
        time.sleep(get_config("riddle.sleep_time"))
        bot.manually_post_message(txt, gid)
    return cb


@as_plugin
def riddle(bot, message, log):
    if message.raw_args is None or message.raw_args != "me piss":
        return RollbotResponse(message, failure=RollbotFailure.INVALID_ARGUMENTS)

    sesh = HTMLSession()
    n = 0
    riddle = None
    while riddle is None:
        rnum, riddle = get_riddle(sesh, log)
        n += 1
        if n > 20:
            return RollbotResponse(message, failure=RollbotFailure.SERVICE_DOWN)

    threading.Thread(target=get_answer_cb(bot, message, riddle)).start()

    return "\n\n".join((
        f"Riddle #{rnum}",
        riddle[0],
        f"I'll post the response in about {get_config('riddle.sleep_time')} seconds!",
    ))


@as_plugin
def fortune():
    sesh = HTMLSession()
    r = sesh.get("http://www.fortunecookiemessage.com/")
    quote = r.html.xpath("//div[contains(@class, 'quote')]/a")[0].text
    learn, lotto = [x.split(":", 1)[-1].strip() for x in r.html.xpath("//div[contains(@class, 'bottom-message')]")[0].text.split("\n")[:2]]

    return f"'{quote}'\nLucky Numbers: {lotto}\nLearn Chinese: {learn}"