Goon Jabber Bot

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
    SleekXMPP: The Sleek XMPP Library
    Copyright (C) 2010  Nathanael C. Fritz
    This file is part of SleekXMPP.

    See the file LICENSE for copying permission.
"""
import os
import sys
import traceback

import pprint
import json
from xml.dom.minidom import parseString
import urllib2 as _urllib # So you can change version without changing code.

import datetime
import logging
import getpass
import random

from optparse import OptionParser

#import MySQLdb as mdb
import shlex # For its cool splitting thingy... http://stackoverflow.com/a/79985/211081
import requests # for its super cool http shit,

import sleekxmpp



# Python versions before 3.0 do not use UTF-8 encoding
# by default. To ensure that Unicode is handled properly
# throughout SleekXMPP, we will set the default encoding
# ourselves to UTF-8.
if sys.version_info < (3, 0):
    #from sleekxmpp.util.misc_ops import setdefaultencoding
    sys.setdefaultencoding('utf8')
else:
    raw_input = input

def restart_program():
    """Restarts the current program.
    Note: this function does not return. Any cleanup action (like
    saving data) must be done before calling this function."""
    python = sys.executable
    os.execl(python, python, * sys.argv)

def millify(n): # http://stackoverflow.com/a/3155023/211081 - stolen
    millnames=['','K','M','B','T']
    millidx=max(0,min(len(millnames)-1,
                      int(math.floor(math.log10(abs(n))/3.0))))
    return '%.0f %s'%(n/10**(3*millidx),millnames[millidx])

class cache_result:
    def __init__ (self, f):
        self.f = f
        self.mem = {}
    def __call__ (self, *args, **kwargs):
        if (args, str(kwargs)) in self.mem:
            return self.mem[args, str(kwargs)]
        else:
            tmp = self.f(*args, **kwargs)
            self.mem[args, str(kwargs)] = tmp
            return tmp


class IntelBot(sleekxmpp.ClientXMPP):

    """
    A simple SleekXMPP bot that will greets those
    who enter the room, and acknowledge any messages
    that mentions the bot's nickname.
    """

    def __init__(self, jid, password, room, nick):
        sleekxmpp.ClientXMPP.__init__(self, jid, password)

        version = 0.1 # set version

        # Do some work to work out if we included a server in the room name and default to @conference.server
        # Will also need to work out if the server was included on the nickname... It should've been but what the hell.
        self.rooms = []

        for v in room.split(','):
            if v.find('@') is -1:
                self.rooms.append("%s@conference.%s" % (v, jid[(jid.find('@')+1):]))
            else:
                self.rooms.append(v)

        self.nick = nick
        self.lastRoom = None

        # Timer to use for ticks
        self.tickTock = 5
        self.timer = True
        self.SCHEDULE_SET = False

        # Some basic session to keep
        self.conn = requests.Session()

        # define some admins
        self.admins = ['psykzz','midge mo\'yb']

        # Database connections
        '''
        self.db_host = 'psykzz.co.uk'
        self.db_user = 'daemo_bot'
        self.db_pass = 'jbQwe123'
        self.db_name = 'yapeal'

        self.conn = mdb.connect(self.db_host, self.db_user,
        self.db_pass, self.db_name)
        self.db = self.conn.cursor(mdb.cursors.DictCursor)
        '''

        # The session_start event will be triggered when
        # the bot establishes its connection with the server
        # and the XML streams are ready for use. We want to
        # listen for this event so that we we can initialize
        # our roster.
        self.add_event_handler("session_start", self.onStart)

        # The groupchat_message event is triggered whenever a message
        # stanza is received from any chat room. If you also also
        # register a handler for the 'message' event, MUC messages
        # will be processed by both handlers.
        self.add_event_handler("groupchat_message", self.on_muc_message)

        # What to do when we call close.
        self.add_event_handler("disconnected", self.onDisconnect)

        # The groupchat_presence event is triggered whenever a
        # presence stanza is received from any chat room, including
        # any presences you send yourself. To limit event handling
        # to a single room, use the events muc::room@server::presence,
        # muc::room@server::got_online, or muc::room@server::got_offline.
        for room in self.rooms:
            self.add_event_handler("muc::%s::got_online" % room,
                               self.on_muc_presence)

    def onStart(self, event):
        """
        Process the session_start event.
        """

        # Before we start the process we need to register the scheduler...
        if not self.SCHEDULE_SET:
            self.schedule('onTickTimer',                       # unique name for the timer
                                self.tickTock,             # seconds to delay before firing
                                self.onTick,               # function to execute
                                #args=('xmpp'),
                                repeat=True)               # make the event happen every X seconds
        self.SCHEDULE_SET = True

        for room in self.rooms:
            self.plugin['xep_0045'].joinMUC(room,
                self.nick,
                wait=True)

    def onDisconnect(self, event):
        # Close the database connection.
        self.conn.close()

    def onTick(self):
        try:
            # End if we shouldn't be running.
            if not self.timer:
                return

            ## DO SOMETHING!
            #print "Test Tock"
            #self.ek_getNewKills()
        except:
            print "fuck - TickTock"

    def on_muc_message(self, msg):
        """
        Process incoming message stanzas from any chat room. Be aware
        that if you also have any handlers for the 'message' event,
        message stanzas may be processed by both handlers, so check
        the 'type' attribute when using a 'message' event handler.

        Whenever the bot's nickname is mentioned, respond to
        the message.

        IMPORTANT: Always check that a message is not from yourself,
                   otherwise you will create an infinite loop responding
                   to your own messages.

        This handler will reply to messages that mention
        the bot's nickname.

        Arguments:
            msg -- The received message stanza. See the documentation
                   for stanza objects and the Message stanza to see
                   how it may be used.
        """



        # Set the last room used to know where to reply to
        self.lastRoom = msg['from'].bare

        # Check for any commamds.
        ###self._parse_cmds(msg)

        try: # Is someone trying to get the bot to send a message? If so, do fuck all... cheeky cunts.
            sender = str(msg['from']).lower()
            sender = sender[sender.index('.com/')+5:]

            if self.nick.lower() == sender:
                return
        except:
            return

        try:
            self._parse_cmds(msg)
        except:
            print '>>> traceback <<<'
            traceback.print_exc()
            print '>>> end of traceback <<<'
            self._send_message("Nope, there is a bug with that.")

        if msg['mucnick'] != self.nick and self.nick in msg['body']:
            self._send_message("I heard that, %s." % msg['mucnick'])

    def on_muc_presence(self, presence):
        """
        Process a presence stanza from a chat room. In this case,
        presences from users that have just come online are
        handled by sending a welcome message that includes
        the user's nickname and role in the room.

        Arguments:
            presence -- The received presence stanza. See the
                        documentation for the Presence stanza
                        to see how else it may be used.
        """
        if presence['muc']['nick'] != self.nick:
            print "%s - [%s] %s" % ( presence['muc']['room'][:presence['muc']['room'].find("@")], presence['muc']['role'][:1], presence['muc']['nick'] )

    def _parse_cmds(self, message):
        """
        With an unproccessed message check if any of the matching commands
        match the start of the string and call the matching function.

        Arguements:
            message -- The unproccessed message.
                        Example "evedata psykzz | evedata 'midge moyb'"
        """

        # Define commands
        self.cmds = ["help","version","restart","joinchannel","leavechannel","intel","translate","pickone"]

        # Parse cmd
        _tmp =  message['body'].lower().split(' ',1)
        cmd = _tmp[0]

        # Check if the cmd exists...
        if not cmd in self.cmds:
            return # Return because it doesnt... eg do nothing...

        # ReParse cmd
        cmd = "cmd_%s" % (cmd)

        if len(_tmp) < 2:
            args = "no arguments"
        else:
            args = _tmp[1]

        # Parse arguments
        ''' Testing attempts.
        arguments = shlex.split(
            _urllib.quote(
                args.encode("utf8")
            )
        )
        arguments = ["%s" % _urllib.unquote(arg).decode("utf8") for arg in arguments]
        #arguments = shlex.split(message.replace('\'','^').split(" ",1)[1])
        #arguments = ["%s" % arg.replace('^','\'') for arg in arguments]
        '''
        arguments = args
        #print arguments.decode("utf8","ignore")
        # Check if the function exists, is callable and call it.
        if hasattr(self, cmd) and callable(getattr(self, cmd)):
            #print u"Function called [%s] with (%s)" % (cmd,arguments.decode("utf8","ignore"))
            getattr(self, cmd)(arguments,message)
        else:
            #print "Function called [%s] with (%s)" % (cmd,arguments.decode("utf8","ignore"))
            print "Erorr! Function [%s] doesn't exist... shit." % (cmd)


    '''
    Parse Cmd shit.
    '''

    def cmd_help(self, arguments, originalMessage):
        self._send_message("PLEH!. Available cmds:")
        self._send_message("- %s" % (', '.join(self.cmds)))
        return

    def cmd_version(self, arguments, originalMessage):
        self._send_message("Version: %s. By PsyKzz" % self.version)
        return

    def cmd_listchannels(self, arguments, originalMessage):
        if not self.isAdmin(originalMessage):
            self._send_message("No")
            return

        originalMessage.reply((", ".join(self.rooms))).send()
        #print (", ".join(self.rooms))
        return

    def cmd_restart(self, _meh, originalMessage):
        if not self.isAdmin(originalMessage):
            self._send_message("No")
            return
        self.disconnect(wait=True)

    def cmd_joinchannel(self, args, originalMessage):
        if not self.isAdmin(originalMessage):
            self._send_message("No")
            return

        args = "%s@conference.goonfleet.com" % (args)
        self.plugin['xep_0045'].joinMUC(args,
                    self.nick,
                    # If a room password is needed, use:
                    # password=the_room_password,
                    wait=True)

    def cmd_leavechannel(self, args, originalMessage):
        if not self.isAdmin(originalMessage):
            self._send_message("No")
            return

        if args == "no arguments":
            args = self.lastRoom
        args = "%s@conference.goonfleet.com" % (args)
        self.plugin['xep_0045'].leaveMUC(args,
                    self.nick)

    def cmd_keyinfo(self, charName, originalMessage, returnString=False):
        if not self.isAdmin(originalMessage):
            if not returnString:
                self._send_message("No")
            return

        charid = self.eve_getID(charName)
        #ret = self.db_getKeyInfo(charid)
        if not returnString:
            self._send_message("Keys on file - %s" % ret)
        else:
            return ret

    def cmd_pickone(self, text, originalMessage):
        _tmp = text.split(' or ')
        rand = random.randint(0, len(_tmp) )
        choice = _tmp[random.randint(0,len(_tmp)-1)] # Inclusive so remove one.
        self._send_message(choice)

    def cmd_translate(self, text, originalMessage):
        url = "http://translate.google.com/translate_a/t"
        payload = {'client': 't', 'sl' : 'auto', 'tl': 'en','ie': 'UTF-8','multires': '1','text': text}
        r = requests.get(url, params=payload)
        self._send_message("- %s" % r.text[r.text.index("[\"")+2:r.text.index("\",\"")])

    def cmd_translate2(self, text, originalMessage):
        url = "http://translate.google.com/translate_a/t"
        payload = {'client': 't', 'sl' : 'auto', 'tl': 'en','ie': 'UTF-8','multires': '1','text': text}
        r = requests.get(url, params=payload)
        self._send_message("- %s" % r.text[r.text.index("[\"")+2:r.text.index("\",\"")])

    def cmd_bitcoin(self, convert, originalMessage):
        url = "http://btcrate.com/convert"
        payload = {'from': 'btc', 'to': convert.lower(), 'exch': 'mtgox', 'conv': 'google', 'amount': 1}

        r = requests.get(url, params=payload)
        self._send_message("- 1 Bitcoin : %s %s" % (str(round(float(r.text[r.text.index(": \"")+3:r.text.index("\"}")]),2)), convert))

    def cmd_intel(self, charName, originalMessage):
        # Get character ID and then data.
        try:
            charID = int(self.eve_getID(charName))
        except Exception:
            charID = 0

        if charID == 0:
            self._send_message("Character ( %s ) Not Found." % charName)
            return

        charData = self.eve_getCharData(charID)
        self._send_message("\n- %s" % (charData))
        #extraData = self.cmd_keyinfo(charName,originalMessage,True)
        #self._send_message("\n- %s\n- API Keys in DB (%s)" % (charData,extraData))


    def cmd_dbintel(self, charName, originalMessage):
        if not self.isAdmin(originalMessage):
            self._send_message("No")
            return
        # Get charID
        try:
            charID = int(self.eve_getID(charName))
        except Exception:
            charID = 0

        if charID == 0:
            self._send_message("Character ( %s ) Not Found." % charName)
            return

        # Get stuff
        #rows = self.db_getCharacterInfo(charID)
        data = "\n- %s [%s] <%s>\n- Skillpoints: %s\n- Balance: %sISK\n- DoB: %s" % \
            (rows['name'],rows['corporationName'],rows['allianceName'],rows['totalSP'],rows['balance'],rows['DoB'],)

        self._send_message(data)


    '''
    Data base shit.
    '''
'''
    def db_getKeyInfo(self, charID):

        # Build the query, Assign params and execute
        self.db.execute("SELECT * \
FROM v1_accountKeyBridge as keyb, \
v1_utilRegisteredKey as utilk \
WHERE keyb.keyID = utilk.keyID \
AND keyb.characterID = %s \
AND utilk.isActive = 1",
                        (charID))

        # Get number of rows.
        numrows = int(self.db.rowcount)
        return str(numrows)
        # loop through results and print.
        ##for i in range(numrows):
        ##    row = self.db.fetchone()
        ##    print row[0], row[1]

    def db_getCharacterInfo(self, charID):

        # Build the query, Assign params and execute
        self.db.execute("SELECT \
cs.name, cs.corporationName, cs.allianceName, cs.balance, cs.DoB,  sum(sp.skillpoints) as totalSP \
FROM v1_charCharacterSheet as cs, \
v1_charSkills as sp \
WHERE cs.characterID = %s \
AND cs.characterID = sp.ownerID \
AND sp.published = 1",
                (int(charID)))

        return self.db.fetchone()
'''
    '''
    EVE Api shit.
    '''
    def eve_getCharData(self, charID):

        url = "https://api.eveonline.com/eve/characterInfo.xml.aspx"
        payload = {'characterID': charID}

        r = requests.get(url, params=payload)
        dom = parseString(r.text)

        alliance = dom.getElementsByTagName('alliance')
        if not alliance:
            alliance = "No Alliance"
        else:
            alliance = dom.getElementsByTagName('alliance') [0].firstChild.nodeValue

        result = "%s [%s] <%s> \n- Previous Corps : " % (
                dom.getElementsByTagName('characterName')[0].firstChild.nodeValue,
                dom.getElementsByTagName('corporation')[0].firstChild.nodeValue,
                alliance
            )

        # In progress
        # Grab corporation IDs and then remove duplicates
        corpIDs = [node.getAttributeNode('corporationID').nodeValue for node in dom.getElementsByTagName('row')]
        corpIDs = list(set(corpIDs))

        # Convert ID to Name (Old way)
        #corpNames = [self.eve_getName(name) for name in corpIDs]


        # Before converting... get a list of alliances from the corp ids.

        ##for _id in corpIDs:
        ##    allyName = self.eW_getAllianceName(_id)


        # Convert ID to Name (New way)
        corpNames = self.eve_getNames(corpIDs)

        result = "%s %s" % ( result, ", ".join(corpNames) )

        # Add zKillboard
        result = "%s\n- https://zkillboard.com/character/%s/" % ( result, charID )
        # get zBoard last kills
        #try:
        zBoardActivity = self.zkb_getKbActivity(charID)
        #finally:
            #zBoardKills = "Error getting Kills"
        result = "%s - [%s] Total Kills / Losses (7 Days)" % ( result, zBoardActivity )

        # Add eveWho
        result = "%s\n- http://evewho.com/pilot/%s" % ( result, dom.getElementsByTagName('characterName')[0].firstChild.nodeValue.replace(" ","%20") )

        return result
    #@cache_result
    def eve_getID(self, eveName):

        url = "https://api.eveonline.com/eve/characterID.xml.aspx"
        payload = {'names': eveName}
        r = requests.get(url, params=payload)
        dom = parseString(r.text)

        xmlTag = dom.getElementsByTagName('row')[0]
        atr = xmlTag.getAttributeNode('characterID')
        return atr.nodeValue

    #@cache_result
    def eve_getName(self, eveID):
        url = "https://api.eveonline.com/eve/characterName.xml.aspx"
        payload = {'ids': eveID}
        r = requests.get(url, params=payload)
        dom = parseString(r.text)

        xmlTag = dom.getElementsByTagName('row')[0]
        atr = xmlTag.getAttributeNode('name')
        return atr.nodeValue

    #@cache_result
    def eve_getNames(self, eveID):
        url = "https://api.eveonline.com/eve/characterName.xml.aspx"
        payload = {'ids': ",".join(eveID)}

        req = requests.Request('GET', 'https://api.eveonline.com/eve/characterName.xml.aspx', params=payload)
        r = req.prepare()
        r.url = r.url.replace("%2C",",")
        s = requests.Session()
        rec = s.send(r)
        #r = requests.get(url, params=payload)
        dom = parseString(rec.text)
        xmlTag = dom.getElementsByTagName('row')
        ret = [tag.getAttributeNode('name').nodeValue for tag in xmlTag]
        return ret

    '''
    EVE WHO shit.
    '''

        # The fuck is this...
    def eW_getAllianceName(self, corpID):
        url = "http://evewho.com/api.php"
        payload = {'type': 'corplist','id': corpID}
        r = requests.get(url, params=payload)

        #Hack out the alliance id, even tho we are working with json
        print r.text
        _tmp = json.loads(str(r.text))
        print _tmp
        pass


    '''
    zKillboard Api shit.
    '''

    def zkb_getNumKills(self, charID, monthsSince=1):

        dateFrom = "%s0000" % (datetime.date.today() - datetime.timedelta(monthsSince*365/12)).isoformat().replace('-','') # 1 month ago
        self._send_message( "------ %s ------" % dateFrom)
        url = "http://zkillboard.com/kills/api/characterID/%s/pastSeconds/604800/orderDirection/asc/no-attackers/no-items/api-only/" % (charID)
        r = requests.get(url)

        tmp = json.loads(str(r.text))
        return len(tmp)

    def zkb_getNumLosses(self, charID, monthsSince=1):

        dateFrom = "%s0000" % (datetime.date.today() - datetime.timedelta(monthsSince*365/12)).isoformat().replace('-','') # 1 month ago
        url = "http://zkillboard.com/api/losses/characterID/%s/pastSeconds/604800/orderDirection/asc/no-attackers/no-items/api-only/" % (charID)
        r = requests.get(url)

        tmp = json.loads(str(r.text))
        return len(tmp)

    def zkb_getKbActivity(self, charID, monthsSince=1):

        dateFrom = "%s0000" % (datetime.date.today() - datetime.timedelta(monthsSince*365/12)).isoformat().replace('-','') # 1 month ago
        url = "http://zkillboard.com/api/characterID/%s/pastSeconds/604800/orderDirection/asc/no-attackers/no-items/api-only/" % (charID)
        r = requests.get(url)

        tmp = json.loads(str(r.text))
        return len(tmp)

    def isAdmin(self, user):
        _tmp = str(user['from']).lower()
        for admin in self.admins:
            if _tmp[_tmp.index('com/')+4:] == admin:
                return True

        return False

    def _send_message(self, message, recepient=False):
        """
        Custom function to send a message to current channel
        Literally used to automate my tasks.
        """
        if message is None:
            return

        if not recepient:
            sendTo = self.lastRoom
        else:
            sendTo = recepient

        self.send_message(mto=sendTo,
            mbody=message,
            mtype='groupchat')


if __name__ == '__main__':
    # Setup the command line arguments.
    optp = OptionParser()

    # Output verbosity options.
    optp.add_option('-q', '--quiet', help='set logging to ERROR',
                    action='store_const', dest='loglevel',
                    const=logging.ERROR, default=logging.INFO)
    optp.add_option('-d', '--debug', help='set logging to DEBUG',
                    action='store_const', dest='loglevel',
                    const=logging.DEBUG, default=logging.INFO)
    optp.add_option('-v', '--verbose', help='set logging to COMM',
                    action='store_const', dest='loglevel',
                    const=5, default=logging.INFO)

    # JID and password options.
    optp.add_option("-j", "--jid", dest="jid",
                    help="JID to use")
    optp.add_option("-p", "--password", dest="password",
                    help="password to use")
    optp.add_option("-r", "--room", dest="room",
                    help="MUC room to join")
    optp.add_option("-n", "--nick", dest="nick",
                    help="MUC nickname")

    opts, args = optp.parse_args()

    # Setup logging.
    logging.basicConfig(level=opts.loglevel,
                        format='%(levelname)-8s %(message)s')

    if opts.jid is None:
        opts.jid = raw_input("Username: ")
    if opts.password is None:
        opts.password = getpass.getpass("Password: ")
    if opts.room is None:
        opts.room = raw_input("MUC room (multiple supported with commas): ")
    if opts.nick is None:
        opts.nick = raw_input("MUC nickname: ")

    # Do some work to work out if we included a server in the room name and default to @conference.server
    # Will also need to work out if the server was included on the nickname... It should've been but what the hell.
    if not opts.jid.find('@') is -1:
        if opts.room.find('@') is -1:
            opts.room = "%s@conference.%s" % (opts.room, opts.jid[(opts.jid.find('@')+1):])

    # Setup the IntelBot and register plugins. Note that while plugins may
    # have interdependencies, the order in which you register them does
    # not matter.
    xmpp = IntelBot(opts.jid, opts.password, opts.room, opts.nick)
    xmpp.register_plugin('xep_0030') # Service Discovery
    xmpp.register_plugin('xep_0045') # Multi-User Chat
    xmpp.register_plugin('xep_0199') # XMPP Ping

    # Connect to the XMPP server and start processing XMPP stanzas.
    if xmpp.connect():
        # If you do not have the dnspython library installed, you will need
        # to manually specify the name of the server if it does not match
        # the one in the JID. For example, to use Google Talk you would
        # need to use:
        #
        # if xmpp.connect(('talk.google.com', 5222)):
        #     ...
        xmpp.process(block=True)
        print("Done")
    else:
        print("Unable to connect.")