root/xchat/xdccq.py

#!/usr/bin/env python

# written by jmoiron, jmoiron.net
# licensed under the GNU GPL v2
# a copy of the license can be found:
#   http://www.gnu.org/licenses/gpl.txt

import xchat
import sys
import re
from exceptions import ValueError

__module_name__ = "xdccq"
__module_version__ = "0.4"
__module_description__ = "Xdcc Queue"

# to enable some debugging output, set to True
__debugging__ = False

if __debugging__:
    import traceback

def print_debug(string):
    global __debugging__
    if __debugging__: 
        string = str(string)
        print "\00302" + string + "\003"

def print_error(string):
    string = str(string)
    print "\00304" + string + "\003"

def print_help(string):
    string = str(string)
    print "\00302" + string + "\003"

def print_success(string):
    string = str(string)
    print "\00303" + string + "\003"

def echo(string):
    string = str(string)
    print "\00300" + string + "\003"

print_info = print_success

cmd_help = {}
gen_help = [
    "\037[help, ?]\037 <cmd>    - prints this help or detailed help for 'cmd'\n",
    "\037[ls, list]\037         - lists files in queue\n",
    "\037get\037 [bot] [#, #-#] - adds 'send' cmds to queue\n",
    "\037rm\037 [bot] <#, #-#>  - removes 'send' cmds from queue\n",
    "\037[cancel,stop]\037      - cancles current transfer",
]
cmd_help['msg'] = "   " + "   ".join(gen_help)
cmd_help['?'] = """\n\002/xdccq\002 \037[help, ?]\037 <cmd>\n   \00303By itself, help/? prints out the available commands along withshort descriptions of each.  If you supply a command \002cmd\002, a longer help description of that command and its usage is given.\003"""
cmd_help['ls'] = """\n\002/xdccq\002 \037[ls, list]\037\n   \00303Lists the file gets currently in the queue.\00303"""
cmd_help['get'] = """\n\002/xdccq\002 \037get\037 [bot] [#, #-#]\n   \00303Adds commands to the local queue "/ctcp \037bot\037 xdcc send \037#\037".  A number ["7"], a number range ["1-5"], or a comma separated list ["7, 8-11, 15"] may be given.\003"""
cmd_help['rm'] = """\n\003/xdccq\002 \037rm\037 [bot] <#, #-#>\n   \00303Removes commands from local queue.  If only \037bot\037 is supplied, it removes all commands dealing with \037bot\037.  If it is supplied with a number or number range, any commands in that range are removed.   If you had numbers "8, 9, 11, 12, 14" queued, and removed "8-14", it would remove all of those package numbers for that bot without resulting in an error for the missing "10" and "13".\003\n   \00303New in 0.3: rm will now also call 'cancel' if the active transfer is from [bot] and the optional range argument was not given.\003"""
cmd_help['cancel'] = """\n\003/xdccq\002 \037cancel\037\n   \00303If a DCC handled by xdccq is currently active, it cancels that transfer and starts the next one on the queue.  If the 'active' transfer is remotely queued, xdccq attempts to unqueue it by issuing a '/msg <bot> xdcc remove' command.  If this command fails, this transfer will continue to remotely pend; take notice!\003"""

# converts a string like "3,5,7-9,14" into a list
# raises InvalidRange
# raises InvalidInteger
def numToList(string):
    ret = []
    numsplit = string.split(",")
    # the following code makes nums into a list of all integers
    for n in numsplit:
        nr = n.split('-')
        # handle the case of a single number
        if len(nr) == 1:
            try: ret.append(int(n))
            except: raise ValueError("number")
        # handle the case of a range
        elif len(nr) == 2:
            try:
                low = int(nr[0])
                high = int(nr[1]) + 1
                if low > high: raise ValueError("number")
                ret += range(low, high)
            except ValueError: raise ValueError("number")
        else: raise ValueError("range")
    return ret

class Command:
    def __init__(self, bot, num):
        self.bot = str(bot)
        self.num = int(num)
        self.channel = xchat.get_info("channel")
        self.network = xchat.get_info("network")
        self.context = xchat.get_context()
        self.file = ""
        self.retries = 0
        self.queue_position = 0
        self.queued = False
        self.transfering = False
        self.dead = False
        self.s = "ctcp %s xdcc send #%d" % (str(bot), int(num))

    def __str__(self):
        return "#%d on %s (%s)" % (self.num, self.bot, self.channel)

class Queue:
    def __init__(self):
        self.data = []
    
    def put(self, item):
        self.data.append(item)
    
    def get(self):
        tmp = self.data[0]
        del self.data[0]
        return tmp
    
    def __getitem__(self, key):
        return self.data[key]
    
    def __delitem__(self, key):
        del self.data[key]
    
    def __len__(self):
        return len(self.data)

    def __iter__(self):
        return self.data.__iter__()

class cmdQueue(Queue):
    """Utility functions for dealing with a queue of Commands"""

    def getBotSet(self, botname, r=[]):
        packs = [cmd for cmd in self.data if cmd.bot == botname]
        # if supplied a range, cut list to those in the range
        if r: packs = [cmd for cmd in packs if cmd.num in r]
        return packs

    def removeBot(self, botname, r=[]):
        packs = self.getBotSet(botname, r)
        for cmd in packs:
            self.data.remove(cmd)
        # this might be better as the actual list someday?
        return len(packs)
        

CommandQueue = cmdQueue()
Active = False
Watchdog = False

# help, ?  -- print out explanations of commands
def help(a):
    global cmd_help
    if len(a) == 3:
        if cmd_help.has_key(a[2]):
            print_help(cmd_help[a[2]])
            return
    usage()
    print cmd_help['msg']

# ls, list -- print out a list of files in the queue
def ls(a):
    global CommandQueue, Active
    s = "No files being transfered."
    if Active:
        if Active.transfering: 
            dcc_list = xchat.get_list("dcc")
            cps = "Unknown CPS"
            for dcc_item in dcc_list:
                if str(dcc_item.nick).lower() == Active.bot.lower():
                    cps = "%s CPS" % (str(dcc_item.cps))
            s = "Pack #%d from %s on %s@%s being transfered at %s." % (Active.num, Active.bot, Active.channel, Active.network, cps)
        elif Active.queued: s = "Pack #%d from %s on %s@%s queued remotely [position %d]." % (Active.num, Active.bot, Active.channel, Active.network, Active.queue_position)
    if len(CommandQueue) == 0:
        print_info("No files in the queue.  " + s)
    else:
        print_info(s)
        for f in CommandQueue:
            print_info(f)

# get -- add commands to the queue
def get(a):
    if len(a) != 4:
        print_error("Error: invalid arguments for get.  /xdccq get [bot] [#, #-#]")
        return
    bot = str(a[2])
    try: nums = numToList(a[3]) 
    except ValueError, exc:
        # exc.args[0] is in the set ['number', 'range']
        print_error("Error: %s contains invalid %s." % (a[3], exc.args[0]))
        return
    print_info("adding packs: %s" % (str(nums)))
    for num in nums:
        CommandQueue.put(Command(bot, num))
    # if we actually added something, try to start this party
    if len(nums):
        run()

# rm -- remove commands from the queue  
def rm(a):
    global CommandQueue, Active
    if len(a) not in (3, 4):
        print_error("Error: invalid arguments for 'rm'.  /xdccq rm [bot] <#, #-#>")
    items_to_delete = []
    if len(a) == 4:
        try: items_to_delete = numToList(a[3])
        except ValueError, exc:
            # exc.args[0] is in the set ['number', 'range']
            print_error("Error: %s contains invalid %s." % (a[3], exc.args[0]))
            return
    # if we removed all from bot, and the bot is currently transfering...
    removed = CommandQueue.removeBot(a[2], items_to_delete)
    if Active and Active.bot == a[2] and len(a) == 3: cancel()
    print_info("deleted %d commands from the queue." % removed)

# we won't be needing a, and it'l clean the call elsewhere
def cancel(a=[]):
    global Active
    if not Active:
        print_error("No currently transfering packages.")
        return
    if Active.transfering:
        s = "dcc close get %s %s" % (Active.bot, Active.file)
        print_debug("/%s" % (s))
        Active.context.command(s)
    elif Active.queued:
        s = "msg %s xdcc remove" % (Active.bot)
        print_debug("/%s" % (s))
        Active.context.command(s)
    Active = False
    run()

USAGE_STR = "Usage: \002/xdccq\002 [cmd] [args], \002/xdccq\002 \037help\037 for commands."

def usage():
    print USAGE_STR

def dispatch(argv, arg_to_eol, c):
    print_debug(argv)
    echo("/" + str(arg_to_eol[0]))
    try:
        { 
        "help"    : help,
        "?"       : help,
        "ls"      : ls,
        "list"    : ls,
        "get"     : get,
        "rm"      : rm,
        "cancel"  : cancel,
        "stop"    : cancel,
    }[argv[1]](argv)
    except:
        if __debugging__: traceback.print_exc(sys.stdout)
        usage()
    return xchat.EAT_XCHAT

def retryTransfer():
    if Active.retries < 3:
        Active.context.command(Active.s)
        Active.retries += 1
        return True
    return False

# watchdog callback
def transferCheck(data):
    global Active
    # if active has already finished, stop the timer
    if not Active: 
        return False
    elif Active.transfering or Active.queued: 
        return False 
    elif not Active.dead:
        Active.dead = True
        return True
    else:
        ret = retryTransfer()
        if ret:
            print_error("Previous attempt to get file (#%d from %s) seems to have failed.  Repeating (%d of 3 retries)" % (Active.num, Active.bot, Active.retries))
        return ret
        
def run():
    global CommandQueue, Active, Watchdog
    if Active: return
    if len(CommandQueue):
        cmd = CommandQueue.get()
        Active = cmd
        print_debug("/%s on %s@%s" % (cmd.s, cmd.channel, cmd.network))
        Active.context.command(Active.s)
        # set up a watchdog every 45 seconds
        Watchdog = xchat.hook_timer(45000, transferCheck)

""" some other messages are possible: 
** Closing Connection: Unable to transfer data (Broken pipe)
** Closing Connection: DCC Timeout (180 Sec Timeout)
"""
def notice(split, full, data):
    global CommandQueue, Active
    # ['jonas|srvr', 'Total Offered: 0.3 MB  Total Transferred: 0.30 MB']
    botname = split[0]
    message = split[1]
    if not Active: return
    if Active.queued: return
    if Active.bot == botname:
        re_queued = re.compile(r"queue")
        re_pos = re.compile(r"position [0-9]+")
        if re_queued.search(message.lower()):
            Active.queued = True
            print_info("xdccq detected the active file being placed on a remote queue")
        if re_pos.search(message.lower()):
            try:
                res = re_pos.search(message.lower())
                postr = message[res.start() : res.end()]
                pos = int(postr.split()[1])
                Active.queue_position = pos
            except:
                Active.queue_position = 0
                print_error("an error occured parsing the remote queue position")
        
def dccComplete(split, full, data):
    # ['just4u.txt', '/home/jmoiron/.xchat2/downloads/just4u.txt.1', 'just|4|u', '38373']
    global CommandQueue, Active
    if not Active: return
    #print_debug(str(split))
    if Active.bot == str(split[2]):
        # we assume it was our file
        print_info("Received file \"%s\" from %s on %s@%s." % (Active.file, Active.bot, Active.channel, Active.network))
        Active = False
    run()

def dccConnect(split, full, data):
    global CommandQueue, Active
    if not Active: return
    #print_debug(str(split))
    botname = split[0]
    if Active.bot == str(split[0]):
        print_info("Requested file \"%s\" is being sent." % (split[2]))
        Active.file = split[2]
        Active.queued = False
        Active.transfering = True

def dccStall(split, full, data):
    global CommandQueue, Active
    if not Active: return
    if Active.bot == str(split[2]):
        print_error("Requested file \"%s\" has stalled during transport." % (split[1]))
        ret = retryTransfer()
        if ret:
            print_info("Re-requesting file \"%s\" (%d of 3 retries)" % (split[1], Active.retries))
        else:
            print_error("Retry limit reached for file \"%s\".  Stopping the queue." % (split[1]))

__unhook__ = xchat.hook_command("xdccq", dispatch, help=USAGE_STR)

print_info("XdccQ-TNG loaded successfully")
usage()

noticeHook = xchat.hook_print("Notice", notice, "data")
dccRecvCompleteHook = xchat.hook_print("DCC RECV Complete", dccComplete, "data")
dccRecvConnectHook = xchat.hook_print("DCC RECV Connect", dccConnect, "data")
dccRecvStallHook = xchat.hook_print("DCC Stall", dccStall, "data")