#! /usr/bin/env python3

import irc
import os
import optparse
import asynchat
import socket
import asyncore
from urllib.parse import quote_plus as quote

nobody = '\002[nobody]\002'

class Flagger(asynchat.async_chat):
    """Connection to flagd"""

    def __init__(self, addr, auth):
        asynchat.async_chat.__init__(self)
        self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
        self.connect((addr, 6668))
        self.push(auth + b'\n')
        self.flag = None

    def handle_read(self):
        msg = self.recv(4096)
        raise ValueError("Flagger died: %r" % msg)

    def handle_error(self):
        # If we lose the connection to flagd, nobody can score any
        # points.  Terminate everything.
        asyncore.close_all()
        asynchat.async_chat.handle_error(self)

    def set_flag(self, team):
        if team:
            eteam = team.encode('utf-8')
        else:
            eteam = b''
        self.push(eteam + b'\n')
        self.flag = team


class Kevin(irc.Bot):
    def __init__(self, host, flagger, tokens, victims):
        irc.Bot.__init__(self, host, ['kevin'], 'Kevin', ['#kevin'])
        self.flagger = flagger
        self.tokens = tokens
        self.victims = victims
        self.affiliation = {}

    def cmd_join(self, sender, forum, addl):
        if sender.name() in self.nicks:
            self.tell_flag(forum)

    def err(self, exception):
        """Save the traceback for later inspection"""
        irc.Bot.err(self, exception)
        t,v,tb = exception
        info = []
        while 1:
            info.append('%s:%d(%s)' % (os.path.basename(tb.tb_frame.f_code.co_filename),
                                       tb.tb_lineno,
                                       tb.tb_frame.f_code.co_name))
            tb = tb.tb_next
            if not tb:
                break
        del tb                          # just to be safe
        infostr = '[' + '] ['.join(info) + ']'
        self.last_tb = '%s %s %s' % (t, v, infostr)
        print(self.last_tb)

    def tell_flag(self, forum):
        forum.msg('%s has the flag.' % (self.flagger.flag or nobody))

    def cmd_privmsg(self, sender, forum, addl):
        text = addl[0]
        if text.startswith('!'):
            parts = text[1:].lower().split(' ', 1)
            cmd = parts[0]
            if len(parts) > 1:
                args = parts[1]
            else:
                args = None
            if cmd.startswith('r'):
                # Register
                who = sender.name()
                if args:
                    self.affiliation[who] = args
                team = self.affiliation.get(who, nobody)
                forum.msg('%s is playing for %s' % (who, team))
            elif cmd.startswith('e'):
                # Embrace
                forum.ctcp('ACTION', 'hugs %s' % sender.name())
            elif cmd.startswith('f'):
                # Flag
                self.tell_flag(forum)
            elif cmd.startswith('h'):
                # Help
                forum.msg('Goal: Obtain a token with social engineering.')
                forum.msg('Commands: !help, !flag, !register [TEAM], !claim TOKEN, !victims, !embrace')
            elif cmd.startswith('c') and args:
                # Claim
                sn = sender.name()
                team = self.affiliation.get(sn)
                token = quote(args, safe='')
                fn = os.path.join(self.tokens, token)
                if not team:
                    forum.msg('%s: register first (!register TEAM).' % sn)
                elif self.flagger.flag == team:
                    forum.msg('%s: Greedy, greedy.' % sn)
                elif not os.path.exists(fn):
                    forum.msg('%s: Token does not exist (possibly already claimed).' % sn)
                else:
                    os.unlink(fn)
                    self.flagger.set_flag(team)
                    self.tell_flag(forum)
            elif cmd.startswith('v'):
                # Victims
                # Open the file each time, so it can change
                try:
                    for line in open(self.victims):
                        forum.msg(line.strip())
                except IOError:
                    forum.msg('There are no victims!')
            elif cmd == 'traceback':
                forum.msg(self.last_tb or 'No traceback')

def main():
    p = optparse.OptionParser()
    p.add_option('-t', '--tokens', dest='tokens', default='./tokens',
                 help='Directory containing tokens')
    p.add_option('-v', '--victims', dest='victims', default='victims.txt',
                 help='File containing victims information')
    p.add_option('-i', '--ircd', dest='ircd', default='localhost',
                 help='IRC server to connect to')
    p.add_option('-f', '--flagd', dest='flagd', default='localhost',
                 help='Flag server to connect to')
    p.add_option('-p', '--password', dest='password',
                 default='kevin:::7db3e44d53d4a466f8facd7b7e9aa2b7',
                 help='Flag server password')
    opts, args = p.parse_args()

    f = Flagger(opts.flagd, opts.password.encode('utf-8'))
    k = Kevin((opts.ircd, 6667), f, opts.tokens, opts.victims)
    irc.run_forever()

if __name__ == '__main__':
    main()