moth/ctf/flagd.py

194 lines
5.3 KiB
Python
Raw Normal View History

#! /usr/bin/env python3
2009-08-31 21:15:13 -06:00
import asyncore
import asynchat
import socket
import functools
import time
import hmac
import optparse
2009-10-01 12:17:03 -06:00
import os
2009-08-21 13:27:17 -06:00
import traceback
2009-10-05 13:33:20 -06:00
from . import teams
from . import points
from . import pointscli
from . import config
key = b'My First Shared Secret (tm)'
def hexdigest(data):
return hmac.new(key, data).hexdigest()
2009-10-01 12:17:03 -06:00
flags_dir = config.get('global', 'flags_dir')
class Flagger(asynchat.async_chat):
"""Use to connect to flagd and submit the current flag holder."""
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
2009-08-31 21:15:13 -06:00
class Submitter(asyncore.dispatcher):
2009-10-05 13:33:20 -06:00
def __init__(self, host='127.0.0.1', port=6667):
2009-08-31 21:15:13 -06:00
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
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 writable(self):
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)
return ret
def handle_write(self):
dgram = self.pending.pop(0)
self.socket.send(dgram)
def handle_read(self):
dgram, peer = self.socket.recvfrom(4096)
try:
2009-08-31 21:15:13 -06:00
id, txt = points.decode_response(dgram)
except ValueError:
# Ignore invalid packets
return
2009-08-31 21:15:13 -06:00
try:
del self.unacked[id]
except KeyError:
pass
if txt != 'OK':
raise ValueError(txt)
def set_flag(self, cat, team):
if not cat:
return
2009-08-31 21:15:13 -06:00
now = int(time.time())
team = team or teams.house
2009-08-31 21:15:13 -06:00
2009-09-01 09:34:15 -06:00
if self.flags.get(cat) != team:
self.flags[cat] = team
self.submit(now, cat, team, 1)
2009-08-31 21:15:13 -06:00
class Listener(asyncore.dispatcher):
2009-10-05 13:33:20 -06:00
def __init__(self, connection_factory, host='', port=6668):
2009-08-31 21:15:13 -06:00
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):
if not self.cat:
return
2009-08-31 21:15:13 -06:00
self.flag = team
2009-10-01 12:17:03 -06:00
self.submitter.set_flag(self.cat, team)
f = open(os.path.join(flags_dir, self.cat), 'w')
if team:
f.write(team)
2009-08-31 21:15:13 -06:00
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)
2009-09-01 09:34:15 -06:00
self.close()
2009-08-31 21:15:13 -06:00
def start():
submitter = Submitter()
server = Listener(functools.partial(FlagServer, submitter))
return (submitter, server)
def main():
p = optparse.OptionParser()
p.add_option('-p', '--genpass', dest='cat', default=None,
help='Generate a password for the given category')
opts, args = p.parse_args()
if opts.cat:
print('%s:::%s' % (opts.cat, hexdigest(opts.cat.encode('utf-8'))))
return
2009-08-31 21:15:13 -06:00
start()
asyncore.loop()
if __name__ == '__main__':
main()