~thrrgilag/partybot

ce99d7416cabe9f0a96cb21dacbdca3a3161826a — Morgan McMillian 2 years ago a54a61d
pnut native partybot with additional features... #4, #10, #11
4 files changed, 430 insertions(+), 0 deletions(-)

A .gitignore
A database.py
A models.py
A partybot-pnut.py
A .gitignore => .gitignore +66 -0
@@ 0,0 1,66 @@
# ---> Python
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class

# C extensions
*.so

# Distribution / packaging
.Python
env/
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
*.egg-info/
.installed.cfg
*.egg

# PyInstaller
#  Usually these files are written by a python script from a template
#  before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec

# Installer logs
pip-log.txt
pip-delete-this-directory.txt

# Unit test / coverage reports
htmlcov/
.tox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*,cover

# Translations
*.mo
*.pot

# Django stuff:
*.log

# Sphinx documentation
docs/_build/

# PyBuilder
target/

# my other cruft
*.yaml
*.db
*.sublime-project
*.sublime-workspace
.vscode/

A database.py => database.py +21 -0
@@ 0,0 1,21 @@
from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy.ext.declarative import declarative_base
import yaml

with open("config.yaml", "rb") as config_file:
    config = yaml.load(config_file)

engine = create_engine(config['SERVICE_DB'])
db_session = scoped_session(sessionmaker(bind=engine))

Base = declarative_base()

Base.query = db_session.query_property()

def init_db():
    # import all modules here that might define models so that
    # they will be registered properly on the metadata.  Otherwise
    # you will have to import them first before calling init_db()
    import models
    Base.metadata.create_all(bind=engine)

A models.py => models.py +25 -0
@@ 0,0 1,25 @@
from sqlalchemy import Column, ForeignKey, Integer, String, Boolean
from database import Base

class Optout(Base):
    __tablename__ = 'optout'
    id = Column(Integer, primary_key=True)
    userid = Column(Integer, unique=True)

class Karma(Base):
    __tablename__ = 'karma'
    id = Column(Integer, primary_key=True)
    userid = Column(Integer)
    chanid = Column(Integer)
    karma = Column(Integer)

class Queue(Base):
    __tablename__ = 'queue'
    id = Column(Integer, primary_key=True)
    msg = Column(String(2048))

class Preferences(Base):
    __tablename__ = 'preferences'
    id = Column(Integer, primary_key=True)
    userid = Column(Integer, unique=True)
    chimpnut = Column(Boolean)

A partybot-pnut.py => partybot-pnut.py +318 -0
@@ 0,0 1,318 @@
import yaml
import requests
import pnutpy
import websocket
import threading
import logging
import time
import json
import random
import re

from database import db_session, init_db
from sqlalchemy import and_
from models import Karma, Optout, Queue, Preferences

_shutdown = threading.Event()
_connected = threading.Event()

def subscribe(connection_id):
    url = f"https://api.pnut.io/v0/channels/{config['CHANNEL']}/messages"
    url += "?connection_id=" + connection_id
    headers = {'Authorization': "Bearer " + config['ACCESS_TOKEN']}
    r = requests.get(url, headers=headers)
    if r.status_code == 200:
        _connected.set()
        send(config['CHANNEL'], "Partybot online and ready! \o/")
    else:
        logger.error(r)

def send(room, text):
    pnutpy.api.create_message(room, data={'text': text})

def echo(room, text):
    logger.debug(text)
    send(room, text)

def help(room):
    reply = "You can upvote or downvote participants by mentioning them with the following symbols\n\n"
    reply += " to vote up: +1, ++, \U0001F44D \n"
    reply += " to vote down: -1, --, \U0001F44E \n"
    reply += "\n"
    reply += "You can use the following commands\n\n"
    reply += "!karma -- show the current karma ratings\n"
    reply += "!optout -- remove yourself from karma ratings\n"
    reply += "!optin -- add yourself to karma ratings\n"
    reply += "!chimpnut -- toggle special ChimPnut alerts\n"
    reply += "!botsnack -- give @partybot a special treat\n"
    send(room, reply)

def botsnack(room):
    replies = [
        "Nom nom nom.",
        "Ooooh",
        "Yummi!",
        "Delightful.",
        "That makes me happy",
        "How kind!",
        "Sweet.",
        "*burp*",
    ]
    send(room, random.choice(replies))

def botdrink(room):
    replies = [
        "Hold my beer.",
        "Ooooh",
        "Delightful.",
        "That makes me happy",
        "How kind!",
        "Sweet.",
        "*burp*",
        "Ah.. Hiccup!"
    ]
    send(room, random.choice(replies))

def optout(msg):
    karma = Karma.query.filter(Karma.userid == msg.user.id).one_or_none()
    if karma:
        db_session.delete(karma)
    entry = Optout.query.filter(Optout.userid == msg.user.id).one_or_none()
    if entry is None:
        entry = Optout(userid=msg.user.id)
        db_session.add(entry)
        db_session.commit()
    reply = "@" + msg.user.username 
    reply += " you have been removed from the karma table"
    send(msg.channel_id, reply)

def optin(msg):
    entry = Optout.query.filter(Optout.userid == msg.user.id).one_or_none()
    if entry:
        db_session.delete(entry)
    reply = "@" + msg.user.username
    reply += " you are able to earn karma"
    send(msg.channel_id, reply)

def upvote(room, user, prefs):
    karma = Karma.query.filter(Karma.userid == user.id).one_or_none()
    if karma is None:
        karma = Karma(userid=user.id, chanid=room, karma=1)
        db_session.add(karma)
    else:
        karma.karma = karma.karma + 1
    db_session.commit()
    if prefs.chimpnut:
        prefix = "/karma "
    else:
        prefix = ""
    reply = prefix + "@" + user.username
    reply += " now has " + str(karma.karma) + " karma in this channel"
    send(room, reply)

def downvote(room, user, prefs):
    karma = Karma.query.filter(Karma.userid == user.id).one_or_none()
    if karma is None:
        karma = Karma(userid=user.id, chanid=room, karma=-1)
        db_session.add(karma)
    else:
        karma.karma = karma.karma - 1
    db_session.commit()
    if prefs.chimpnut:
        prefix = "/karma "
    else:
        prefix = ""
    reply = prefix + "@" + user.username
    reply += " now has " + str(karma.karma) + " karma in this channel"
    send(room, reply)

def karma(room):
    reply = "Karma standings\n\n"
    results = Karma.query.filter(Karma.chanid == room).order_by(Karma.karma.desc()).all()
    for entry in results:
        user, meta = pnutpy.api.get_user(entry.userid)
        reply += user.username + ": " + str(entry.karma) + "\n"
    send(room, reply)

def chimpnut(msg):
    prefs = Preferences.query.filter(Preferences.userid == msg.user.id).one_or_none()
    if prefs is None:
        prefs = Preferences(userid=msg.user.id, chimpnut=False)
        db_session.add(prefs)
    
    if prefs.chimpnut:
        prefs.chimpnut = False
        reply = "@" + msg.user.username + " ChimPnut alert is now disabled"
    else:
        prefs.chimpnut = True
        reply = "/Mac @" + msg.user.username + " ChimPnut alert is now enabled"
    db_session.commit()
    send(msg.channel_id, reply)

def on_command(msg):
    room = msg.channel_id
    args = msg.content.text.split(' ', 1)

    if args[0] == "!help":
        help(msg.channel_id)

    elif args[0] == "!botsnack":
        botsnack(msg.channel_id)
    
    elif args[0] == "!botdrink":
        botdrink(msg.channel_id)
    
    elif args[0] == "!optout":
        optout(msg)
    
    elif args[0] == "!optin":
        optin(msg)
    
    elif args[0] == "!chimpnut":
        chimpnut(msg)

    elif args[0] == "!karma":
        karma(msg.channel_id)

def on_vote(msg, matcher):

    canidate = matcher.group(1)
    vote = matcher.group(2)

    try:
        pnutuser, meta = pnutpy.api.get_user("@" + canidate)
        if msg.user.username == pnutuser.username:
            logger.debug(pnutuser.username)
            logger.debug(canidate)
            reply = "@" + msg.user.username
            reply += " silly human, your karma must be decided by others!"
            send(msg.channel_id, reply)
            return
    
    except pnutpy.errors.PnutMissing:
        reply = "@" + msg.user.username
        reply += " I do not know who that is"
        send(msg.channel_id, reply)
        return

    optout = Optout.query.filter(Optout.userid == pnutuser.id).one_or_none()
    if optout:
        reply = "@" + msg.user.username
        reply += " user has chosen not to receive karma"
        send(msg.channel_id, reply)
        return

    prefs = Preferences.query.filter(Preferences.userid == pnutuser.id).one_or_none()
    if prefs is None:
        prefs = Preferences(userid=pnutuser.id, chimpnut=True)
        db_session.add(prefs)
        db_session.commit()

    logger.debug("-- VOTING " + vote)

    upvotes = [
        "++",
        ":thumbsup:",
        ":+1:",
        "+1",
        "\U0001F44D"
    ]
    downvotes = [
        "--",
        ":thumbsdown:",
        ":-1:",
        "-1",
        "\U0001F44E"
    ]
    if vote in upvotes:
        upvote(msg.channel_id, pnutuser, prefs)

    elif vote in downvotes:
        downvote(msg.channel_id, pnutuser, prefs)

def on_mention(msg):
    # TODO: use for even more magic
    return

def on_message(ws, message):
    logger.debug("on_message: " + message)
    msg = json.loads(message)

    if not _connected.isSet() and 'connection_id' in msg['meta']:
        send(config['CHANNEL'], "...connecting circuits...")
        logger.debug("connection_id: " + msg['meta']['connection_id'])
        subscribe(msg['meta']['connection_id'])
        return
    
    if 'data' in msg:

        if "channel_type" in msg['meta'] and msg['meta']['channel_type'] == "io.pnut.core.chat":

            pmsg = pnutpy.models.Message.from_response_data(msg['data'])

            if 'is_deleted' in msg['meta']:
                return

            vpattern = r"([\w]+)\s?(\-\-|\+\+|\U0001F44D|\U0001F44E|\:thumbsup\:|\:\+1\:|\:thumbsdown\:|\:-1\:|\+1|-1)"
            votes = re.search(vpattern, pmsg.content.text)

            if pmsg.user.username == config['USERNAME']:
                return

            if config['USERNAME'] in [e.text for e in pmsg.content.entities.mentions]:
                on_mention(pmsg)

            elif pmsg.content.text.startswith('!'):
                on_command(pmsg)

            elif votes:
                on_vote(pmsg, votes)

def on_error(ws, error):
    logger.error("on_error: !!! ERROR !!!")
    logger.error(error)
    _shutdown.set()

def on_close(ws):
    send(config['CHANNEL'], "...shutdown initiated...")
    logger.debug("on_close: ### CLOSED ###")
    _shutdown.set()

def on_open(ws):

    def run(*args):
        while not _shutdown.isSet():
            qmsg = Queue.query.one_or_none()
            if qmsg:
                send(config['CHANNEL'], qmsg.msg)
                db_session.delete(qmsg)
                db_session.commit()
            time.sleep(3)
            ws.send(".")
        time.sleep(1)
        ws.close()
        logger.debug("*** terminate ***")
    
    t = threading.Thread(target=run)
    t.start()

if __name__ == "__main__":

    logger = logging.getLogger()
    logging.basicConfig(level=logging.DEBUG)

    with open("config.yaml", "rb") as config_file:
        config = yaml.load(config_file)
    
    init_db()

    pnutpy.api.add_authorization_token(config['ACCESS_TOKEN'])

    ws_url = "wss://stream.pnut.io/v0/user"
    ws_url += "?access_token=" + config['ACCESS_TOKEN']

    # setup the websocket connection
    ws = websocket.WebSocketApp(ws_url, on_message=on_message, 
        on_error=on_error, on_close=on_close)
    ws.on_open = on_open
    ws.run_forever()