From ce99d7416cabe9f0a96cb21dacbdca3a3161826a Mon Sep 17 00:00:00 2001 From: Morgan McMillian Date: Sun, 17 Feb 2019 08:39:23 -0800 Subject: [PATCH] pnut native partybot with additional features... #4, #10, #11 --- .gitignore | 66 ++++++++++ database.py | 21 ++++ models.py | 25 ++++ partybot-pnut.py | 318 +++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 430 insertions(+) create mode 100644 .gitignore create mode 100644 database.py create mode 100644 models.py create mode 100644 partybot-pnut.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c731d5f --- /dev/null +++ b/.gitignore @@ -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/ diff --git a/database.py b/database.py new file mode 100644 index 0000000..c2f9e22 --- /dev/null +++ b/database.py @@ -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) diff --git a/models.py b/models.py new file mode 100644 index 0000000..982de96 --- /dev/null +++ b/models.py @@ -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) diff --git a/partybot-pnut.py b/partybot-pnut.py new file mode 100644 index 0000000..808e75f --- /dev/null +++ b/partybot-pnut.py @@ -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()