# exploit example

class Exploit():

    def execute(self, ip, port, flag_id):
        from multiprocessing import Process, Queue
        from collections import defaultdict
        from twisted.internet.protocol import Protocol, ClientFactory
        from twisted.internet import reactor
        from twisted.protocols.basic import LineReceiver
        from sys import stdout
        import sys
        import re
        import logging

        NUM_FILES_PER_AGENT = 15

        class File(object):
            def __init__(self, name, size, value):
                self.name = name
                self.size = size
                self.value = value
            def __repr__(self):
                return "Name:{0}, Size:{1}, Value:{2}".format(\
                        self.name, self.size, self.value)

        class Solver(object):
            def __init__(self, host, port, room_id):
                self.agents = []
                self.host = host
                self.port = port
                self.room_id = room_id
                self.flag = None
                reactor.connectTCP(self.host, port, AgentFactory(self))
                reactor.connectTCP(self.host, port, AgentFactory(self))
                reactor.connectTCP(self.host, port, AgentFactory(self))

            def add_agent(self, agent):
                self.agents.append(agent)

            def solve_individual(self):
                # Begin communication
                self.agents[0].msg_snowden("hi")
                reactor.callLater(.1, self.tell_agents_to_send)

            def tell_agents_to_send(self):
                for agent in self.agents:
                    agent.send_files()

            def check_if_ready(self):
                received_files = map(lambda agent: agent.received_all_files(), self.agents)
                if all(received_files):
                    self.solve_individual()


        def memoize(f):
            """ Memoization decorator for functions taking one or more arguments. """
            class memodict(dict):
                def __init__(self, f):
                    self.f = f
                def __call__(self, *args):
                    return self[args]
                def __missing__(self, key):
                    ret = self[key] = self.f(*key)
                    return ret
            return memodict(f)

        class Agent(LineReceiver):
            delimiter = "\n"

            def __init__(self, solver):
                self.solver = solver
                self.solver.add_agent(self)
                self.identity = None
                self.files = []
                self.bandwidth = 0
                self.committed = False

            def received_all_files(self):
                return len(self.files) == NUM_FILES_PER_AGENT

            def determine_identity(self, data):
                find_self = re.search('(\w+) has joined', data)
                if find_self:
                    self.identity = find_self.group(1)

            # list -- | Remaining Bandwidth: 16079 KB
            def determine_bandwidth(self, data):
                find_bandwidth = re.search('(\d+)', data)
                if find_bandwidth:
                    self.bandwidth = int(find_bandwidth.group(1))
                    logging.debug("{0} has bandwidth {1}".format(self.identity,
                            self.bandwidth))

            def msg_snowden(self,msg):
                cmd = "/msg E.Snowden {0}".format(msg)
                if msg == "DONE" and not self.committed:
                    logging.debug("{0} types {1}".format(self.identity, cmd))
                    self.sendLine(cmd)
                    self.committed = True
                elif msg != "DONE":
                    self.sendLine("/msg E.Snowden {0}".format(msg))

            def send_files(self):
                value, files_to_send = self.knapsack(self.files, self.bandwidth)
                for myfile in files_to_send:
                    self.send_file(myfile.name, "E.Snowden")
                    self.files.remove(myfile)
                self.bandwidth -= sum(map(lambda sent_file: sent_file.size,
                                    files_to_send))
                logging.debug("Bandwidth for {0} is now {1}".format(self.identity,
                        self.bandwidth))
                logging.debug("Remaining Files for {0} is {1}".format(self.identity,
                        len(self.files)))

                if len(self.files) >= 3:
                    for agent, my_file in zip(self.solver.agents, self.files[:3]):
                        self.send_file(my_file.name, agent.identity)
                    self.msg_snowden("DONE")

            def send_file (self, file_name, agent):
                cmd = "/send {0} {1}".format(agent, file_name)
                logging.debug("{0} types {1}".format(agent, cmd))
                self.sendLine(cmd)

            def show_up_for_work(self):
                self.sendLine("1")

            def enter_room(self):
                self.sendLine(self.solver.room_id)

            def begin_mission(self):
                self.sendLine("/list")

            def received_file(self, line):
                valid_file = re.search('(\S+)\s+(\d+)KB\s+(\d*)', line)
                if valid_file:
                    new_file = File(valid_file.group(1), int(valid_file.group(2)),
                                    int(valid_file.group(3)))
                    self.files.append(new_file)
                    self.solver.check_if_ready()

            # send -- | *Received File: MasteringTheInternet.ppt from Agent2 *
            def receive_file_from_friend(self, line):
                valid_file = re.search('Received File:\s+(\S+)\((\d+)\)', line)
                if valid_file:
                    new_file_name = valid_file.group(1)
                    file_size = int(valid_file.group(2))
                    if file_size < self.bandwidth:
                        self.send_file(new_file_name, "E.Snowden")
                        self.msg_snowden("DONE")

            def lineReceived(self, data):
                if "Enter the number" in data:
                    self.show_up_for_work()
                elif "Enter your room id" in data:
                    self.enter_room()
                elif self.identity is None and "has joined" in data:
                    self.determine_identity(data)
                elif "Everyone has arrived" in data:
                    self.begin_mission()
                elif "Remaining Bandwidth" not in data and "list -- |" in data:
                    self.received_file(data[12:])
                elif "Remaining Bandwidth" in data:
                    self.determine_bandwidth(data)
                elif "Received File" in data:
                    self.receive_file_from_friend(data)
                # think your boss will want this: test
                elif "boss will want" in data:
                    logging.debug(data)
                    flag = re.search('(\S+)$', data)
                    if flag:
                        self.solver.flag = flag.group(1)
                    reactor.stop()

            def knapsack(self, files, maxweight):

                # Return the value of the most valuable subsequence of the first i
                # elements in items whose weights sum to no more than j.
                @memoize
                def bestvalue(i, j):
                    if i == 0: return 0
                    value = files[i - 1].value
                    weight = files[i - 1].size
                    if weight > j:
                        return bestvalue(i - 1, j)
                    else:
                        return max(bestvalue(i - 1, j),
                                bestvalue(i - 1, j - weight) + value)

                j = maxweight
                result = []
                for i in xrange(len(files), 0, -1):
                    if bestvalue(i, j) != bestvalue(i - 1, j):
                        result.append(files[i - 1])
                        j -= files[i - 1].size
                result.reverse()
                return bestvalue(len(files), maxweight), result

        class AgentFactory(ClientFactory):
            def __init__(self, solver):
                self.solver = solver

            def buildProtocol(self, addr):
                return Agent(self.solver)

        def run(host, port, room_id, queue):
            logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.CRITICAL)
            # Connect three agents
            solver = Solver(host, port, room_id)
            reactor.run()
            #return solver.flag
            queue.put(solver.flag)

                
        flag = ''
        error = 0
        error_msg = ''

        try:
            queue = Queue()
            p = Process(target=run, args=(ip, port, flag_id, queue))
            p.start()
            p.join()
            flag = queue.get()
        except Exception as e:
            error = -1 #down
            error_msg = str(e)

        self.flag = flag
        self.error = error
        self.error_msg = error_msg

    def result(self):
        return {'FLAG' : self.flag,
                'ERROR' : self.error,
                'ERROR_MSG' : self.error_msg,
               }