From c2d685005baba36cb33c018b9d594ae43a81fb85 Mon Sep 17 00:00:00 2001 From: Neale Pickett Date: Mon, 31 Aug 2009 21:15:13 -0600 Subject: [PATCH] Convert flagd to use medusa --- flagd.py | 176 ++++++++++++++++++++++++++++++++++++----------------- ubersrv.py | 12 ++++ 2 files changed, 132 insertions(+), 56 deletions(-) create mode 100755 ubersrv.py diff --git a/flagd.py b/flagd.py index 40cf4e8..5692ff7 100755 --- a/flagd.py +++ b/flagd.py @@ -1,8 +1,9 @@ #! /usr/bin/env python3 -import socketserver -import threading -import queue +import asyncore +import asynchat +import socket +import functools import time import hmac import optparse @@ -14,63 +15,129 @@ key = b'My First Shared Secret (tm)' def hexdigest(data): return hmac.new(key, data).hexdigest() -house = 'dirtbags' # House team name -flags = {} -toscore = queue.Queue(50) +class Submitter(asyncore.dispatcher): + def __init__(self, host='localhost', port=6667): + asyncore.dispatcher.__init__(self) + self.create_socket(socket.AF_INET, socket.SOCK_DGRAM) + self.connect((host, port)) + self.pending = [] + self.unacked = {} + self.flags = {} + self.lastupdate = 0 + self.lastretrans = 0 + self.id = 0 -class Submitter(threading.Thread): - def run(self): - self.sock = pointscli.makesock('localhost') - while True: - try: - delay = 60 - (time.time() % 60) - cat, team = toscore.get(True, delay) - self.submit(cat, team) - except queue.Empty: - self.once() + def submit(self, now, cat, team, score): + q = points.encode_request(self.id, now, cat, team, score) + self.id += 1 + self.pending.append(q) + self.unacked[id] = q - def once(self): - global flags - global toscore + def writable(self): + print('writable?') + now = int(time.time()) + if now >= self.lastupdate + 60: + for cat, team in self.flags.items(): + self.submit(now, cat, team, 1) + self.lastupdate = now + if now > self.lastretrans: + for id, q in self.unacked.items(): + self.pending.append(q) + self.lastretrans = now + ret = bool(self.pending) + print(ret) + return ret - for cat, team in flags.items(): - self.submit(cat, team) + def handle_write(self): + dgram = self.pending.pop(0) + self.socket.send(dgram) - def submit(self, cat, team): + def handle_read(self): + dgram, peer = self.socket.recvfrom(4096) try: - pointscli.submit(self.sock, cat, team, 1) - except: - traceback.print_exc() - - -class CategoryHandler(socketserver.StreamRequestHandler): - def handle(self): - global flags - - try: - catpass = self.rfile.readline().strip() - cat, passwd = catpass.split(b':::') - passwd = passwd.decode('utf-8') - if passwd != hexdigest(cat): - self.wfile.write(b'ERROR :Closing Link: Invalid password\n') - return - cat = cat.decode('utf-8') - except ValueError as foo: - self.wfile.write(b'ERROR :Closing Link: Invalid command\n') + id, txt = points.decode_response(dgram) + except ValueError: + # Ignore invalid packets return + try: + del self.unacked[id] + except KeyError: + pass + if txt != 'OK': + raise ValueError(txt) - flags[cat] = house - while True: - team = self.rfile.readline().strip().decode('utf-8') - if not team: - break - flags[cat] = team - toscore.put((cat, team)) # score a point immediately - flags[cat] = house + def set_flag(self, cat, team): + now = int(time.time()) -class MyServer(socketserver.ThreadingTCPServer): - allow_reuse_address = True + team = team or points.house + self.flags[cat] = team + self.submit(now, cat, team, 1) + + +class Listener(asyncore.dispatcher): + def __init__(self, connection_factory, host='localhost', port=6668): + asyncore.dispatcher.__init__(self) + self.create_socket(socket.AF_INET, socket.SOCK_STREAM) + self.set_reuse_addr() + self.bind((host, port)) + self.listen(4) + self.connection_factory = connection_factory + + def handle_accept(self): + conn, addr = self.accept() + self.connection_factory(conn) + + +class FlagServer(asynchat.async_chat): + def __init__(self, submitter, sock): + asynchat.async_chat.__init__(self, sock=sock) + self.set_terminator(b'\n') + self.submitter = submitter + self.flag = None + self.inbuf = [] + self.cat = None + + def err(self, txt): + e = ('ERROR: Closing Link: %s\n' % txt) + self.push(e.encode('utf-8')) + self.close() + + def collect_incoming_data(self, data): + if len(self.inbuf) > 10: + return self.err('max sendq exceeded') + self.inbuf.append(data) + + def set_flag(self, team): + self.flag = team + if self.cat: + self.submitter.set_flag(self.cat, team) + + def found_terminator(self): + data = b''.join(self.inbuf) + self.inbuf = [] + if not self.cat: + try: + cat, passwd = data.split(b':::') + passwd = passwd.decode('utf-8') + if passwd != hexdigest(cat): + return self.err('Invalid password') + self.cat = cat.decode('utf-8') + except ValueError: + return self.err('Invalid command') + self.set_flag(None) + else: + team = data.strip().decode('utf-8') + self.set_flag(team) + + def handle_close(self): + self.set_flag(None) + + +def start(): + submitter = Submitter() + server = Listener(functools.partial(FlagServer, submitter)) + return (submitter, server) def main(): p = optparse.OptionParser() @@ -81,11 +148,8 @@ def main(): print('%s:::%s' % (opts.cat, hexdigest(opts.cat.encode('utf-8')))) return - submitter = Submitter() - submitter.start() - server = MyServer(('', 6668), CategoryHandler) - server.serve_forever() - + start() + asyncore.loop() if __name__ == '__main__': main() diff --git a/ubersrv.py b/ubersrv.py new file mode 100755 index 0000000..830b954 --- /dev/null +++ b/ubersrv.py @@ -0,0 +1,12 @@ +#! /usr/bin/env python3 + +import asyncore +import pointsd +import flagd + +def main(): + pointsd.start() + flagd.start() + asyncore.loop(timeout=30, use_poll=True) + +main()