lepprint_bot.py

#!/usr/bin/env python3
import contextlib
import io
import os
import time

import discord  # pip install discord.py
import lepprint.lepprint as lepprint  # https://gist.github.com/thurask/632cdc049767aaa63492e2bd6039a1ca
import requests  # pip install requests

# https://www.devdungeon.com/content/make-discord-bot-python

# Discord bot token, either environment variable or set here
RAWTOKEN = "0xDEADBEEF"
TOKEN = os.environ.get("LEPPRINT_BOT_TOKEN", RAWTOKEN)
# Paste token, either environment variable or set here
RAWPASTETOKEN = "0xBABACAFE"
PASTETOKEN = os.environ.get("LEPPRINT_PASTE_TOKEN", RAWPASTETOKEN)
# Paste API endpoint
PASTEURL = "https://api.paste.ee/v1/pastes"
# ID for bot posting channel
BOTCHANNELID = "530483875051208704"
# Names for support channels
SUPPORTCHANNELS = ("lastexceptions_here", "general_support", "nsfw_adult_mod_support", "mccc_support", "mcwoohoo_support", "bilingual_support")
# Extensions to care about
FILEEXTS = (".txt", ".log")

HERE = os.path.dirname(os.path.abspath(__file__))

client = discord.Client()
sess = requests.Session()
sess.headers.update({"X-Auth-Token": PASTETOKEN, "Content-Type": "application/json"})


def prep_attachments(attachments):
    """Get attachment URLs for any text files."""
    # filter .docx and other trash
    attlist = [attach["url"] for attach in attachments if attach["filename"].endswith(FILEEXTS)]
    return attlist


def download_attachments(attachments):
    """Pull all attachments."""
    filelist = []
    for attach in attachments:
        lfname = attach.split("/")[-1]
        filepath = os.path.join(HERE, lfname)
        fileext = os.path.splitext(filepath)[-1]
        if os.path.exists(filepath):
            filepath = filepath.replace(fileext, "_{}_{}".format(str(time.time()), fileext))
        filelist.append(filepath)
        req = requests.get(attach, stream=True)
        with open(filepath, "wb") as afile:
            for chunk in req.iter_content(chunk_size=2048):
                if chunk:
                    afile.write(chunk)
    return filelist


def lepprint_handler(filelist):
    """Run pretty printer and collect output."""
    with io.StringIO() as sio:
        with contextlib.redirect_stdout(sio):
            for file in filelist:
                lepprint.main(file, clatest=True, consume=True)
        content = sio.getvalue()
    return content


def msgurl_generator(message):
    """Generate URLs to messages."""
    server = message.server.id
    channel = message.channel.id
    msgid = message.id
    msgurl = "https://discordapp.com/channels/{}/{}/{}".format(server, channel, msgid)
    return msgurl


def prep_header(message):
    """Prepare message headers."""
    dfmt = "%Y-%m-%d %H:%M:%S"
    msgskel = "Posted by {} in #{} at {} UTC"
    msghead = msgskel.format(message.author, message.channel, message.timestamp.strftime(dfmt))
    msgurl = msgurl_generator(message)
    header = "\n".join(("-----", msghead, msgurl, "---"))
    return header


def post_buffer(iobuffer):
    """Paste buffer content and get URL."""
    payload = {"sections": [{"contents": iobuffer}]}
    req = sess.post(PASTEURL, json=payload)
    rjay = req.json()
    pasteurl = rjay["link"]
    return pasteurl


@client.event
async def on_message(message):
    """Scan for messages."""
    # we do not want the bot to reply to itself
    if message.author == client.user:
        return

    # ignore non-support channels
    if message.channel.name not in SUPPORTCHANNELS:
        return

    # handle attachments
    if message.attachments:
        atts = prep_attachments(message.attachments)
        if not atts:  # no .txt attachments
            return
        else:  # at least one .txt
            files = download_attachments(atts)
            iobuffer = lepprint_handler(files)
            pasteurl = post_buffer(iobuffer)
            header = prep_header(message)
            msg = "\n".join((header, pasteurl))
            await client.send_message(client.get_channel(BOTCHANNELID), msg)


@client.event
async def on_ready():
    """Start the party."""
    print('Logged in as')
    print(client.user.name)
    print(client.user.id)
    print('------')


if __name__ == "__main__":
    client.run(TOKEN)