Convert flagd to use medusa

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()

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()