moth/ctf/points.py

222 lines
5.0 KiB
Python
Raw Normal View History

2009-08-20 12:51:47 -06:00
#! /usr/bin/env python3
import socket
import hmac
import struct
import io
2009-09-29 15:36:25 -06:00
import os
2009-10-05 13:33:20 -06:00
from . import teams
from . import config
2009-08-20 12:51:47 -06:00
##
## Authentication
##
2009-08-20 16:20:44 -06:00
key = b'mollusks peck my galloping genitals'
2009-08-20 12:51:47 -06:00
def digest(data):
return hmac.new(key, data).digest()
def sign(data):
return data + digest(data)
def check_sig(data):
base, mac = data[:-16], data[-16:]
if mac == digest(base):
return base
else:
raise ValueError('Invalid message digest')
##
## Marshalling
##
def unpack(fmt, buf):
"""Unpack buf based on fmt, return the rest as a buffer."""
size = struct.calcsize(fmt)
vals = struct.unpack(fmt, buf[:size])
return vals + (buf[size:],)
def packstr(s):
b = bytes(s, 'utf-8')
return struct.pack('!H', len(b)) + b
def unpackstr(b):
l, b = unpack('!H', b)
s, b = b[:l], b[l:]
return str(s, 'utf-8'), b
##
## Request
##
2009-09-01 09:34:15 -06:00
def encode_request(id, when, cat, team, score):
base = (struct.pack('!II', id, when) +
2009-08-20 12:51:47 -06:00
packstr(cat) +
packstr(team) +
struct.pack('!i', score))
return sign(base)
def decode_request(b):
base = check_sig(b)
2009-09-01 09:34:15 -06:00
id, when, base = unpack('!II', base)
2009-08-20 12:51:47 -06:00
cat, base = unpackstr(base)
team, base = unpackstr(base)
score, base = unpack('!i', base)
assert not base
2009-09-01 09:34:15 -06:00
return (id, when, cat, team, score)
2009-08-20 12:51:47 -06:00
##
## Response
##
2009-09-01 09:34:15 -06:00
def encode_response(id, txt):
base = (struct.pack('!I', id) +
2009-08-20 12:51:47 -06:00
packstr(txt))
return sign(base)
def decode_response(b):
base = check_sig(b)
2009-09-01 09:34:15 -06:00
id, base = unpack('!I', base)
2009-08-20 12:51:47 -06:00
txt, base = unpackstr(base)
assert not base
2009-09-01 09:34:15 -06:00
return (id, txt)
2009-08-20 12:51:47 -06:00
##
## Storage
##
2009-08-21 13:27:17 -06:00
def incdict(dict, key, amt=1):
2009-09-01 09:34:15 -06:00
dict[key] = dict.get(key, 0) + amt
2009-08-21 13:27:17 -06:00
2009-08-20 12:51:47 -06:00
class Storage:
2009-10-07 08:31:54 -06:00
def __init__(self, fn=None, fix=False):
2009-09-29 15:36:25 -06:00
if not fn:
fn = config.datafile('scores.dat')
2009-10-01 12:17:03 -06:00
self.teams = set()
2009-08-20 12:51:47 -06:00
self.points_by_cat = {}
2009-08-21 13:27:17 -06:00
self.points_by_cat_team = {}
2009-08-20 12:51:47 -06:00
self.log = []
self.f = io.BytesIO()
# Read stored scores
2009-10-07 08:31:54 -06:00
truncate = False
lastgood = 0
2009-08-20 12:51:47 -06:00
try:
f = open(fn, 'rb')
while True:
l = f.read(4)
if not l:
2009-08-21 13:27:17 -06:00
break
2009-08-20 12:51:47 -06:00
(l,) = struct.unpack('!I', l)
b = f.read(l)
2009-09-01 09:34:15 -06:00
when, score, catlen, teamlen, b = unpack('!IiHH', b)
cat = b[:catlen].decode('utf-8')
team = b[catlen:].decode('utf-8')
req = (when, cat, team, score)
self.add(req, False)
2009-10-07 08:31:54 -06:00
lastgood = f.tell()
2009-08-20 12:51:47 -06:00
f.close()
2009-10-07 08:31:54 -06:00
except struct.error:
if fix:
truncate = True
else:
raise
2009-08-20 12:51:47 -06:00
except IOError:
pass
2009-08-25 17:50:49 -06:00
try:
self.f = open(fn, 'ab')
except IOError:
self.f = None
2009-08-20 12:51:47 -06:00
2009-10-07 08:31:54 -06:00
if truncate:
self.f.seek(lastgood)
self.f.truncate()
2009-09-01 09:34:15 -06:00
def __len__(self):
return len(self.log)
2009-08-20 12:51:47 -06:00
2009-09-01 09:34:15 -06:00
def add(self, req, write=True):
2009-08-20 12:51:47 -06:00
when, cat, team, score = req
2009-10-01 12:17:03 -06:00
self.teams.add(team)
2009-08-21 13:27:17 -06:00
incdict(self.points_by_cat, cat, score)
incdict(self.points_by_cat_team, (cat, team), score)
2009-08-20 12:51:47 -06:00
self.log.append(req)
2009-09-01 09:34:15 -06:00
if write:
cat = cat.encode('utf-8')
team = team.encode('utf-8')
b = (struct.pack('!IiHH', when, score, len(cat), len(team)) +
cat + team)
lb = struct.pack('!I', len(b))
self.f.write(lb)
self.f.write(b)
self.f.flush()
2009-08-20 12:51:47 -06:00
def categories(self):
return sorted(self.points_by_cat)
2009-10-01 12:17:03 -06:00
def get_teams(self):
return sorted(self.teams)
2009-08-20 12:51:47 -06:00
def cat_points(self, cat):
2009-08-21 13:27:17 -06:00
return self.points_by_cat.get(cat, 0)
2009-08-20 12:51:47 -06:00
def team_points(self, team):
2009-10-01 12:17:03 -06:00
points = 0
for cat, tot in self.points_by_cat.items():
team_points = self.team_points_in_cat(cat, team)
points += team_points / float(tot)
return points
2009-08-21 13:27:17 -06:00
def team_points_in_cat(self, cat, team):
return self.points_by_cat_team.get((cat, team), 0)
2009-08-20 12:51:47 -06:00
2009-08-25 17:50:49 -06:00
2009-08-20 12:51:47 -06:00
##
## Testing
##
def test():
import time
import os
now = int(time.time())
req = (now, 'category 5', 'foobers in heat', 43)
assert decode_request(encode_request(*req)) == req
2009-09-01 09:34:15 -06:00
rsp = (now, 'cat6', 'hello world')
2009-08-20 12:51:47 -06:00
assert decode_response(encode_response(*rsp)) == rsp
try:
os.unlink('test.dat')
except OSError:
pass
s = Storage('test.dat')
s.add((now, 'cat1', 'zebras', 20))
s.add((now, 'cat1', 'aardvarks', 10))
s.add((now, 'merf', 'aardvarks', 50))
2009-10-01 12:17:03 -06:00
assert s.get_teams() == ['aardvarks', 'zebras']
2009-08-20 12:51:47 -06:00
assert s.categories() == ['cat1', 'merf']
assert s.cat_points('cat1') == 30
2009-08-21 13:27:17 -06:00
assert s.team_points_in_cat('cat1', 'aardvarks') == 10
assert s.team_points_in_cat('merf', 'zebras') == 0
2009-08-20 12:51:47 -06:00
del s
s = Storage('test.dat')
assert s.teams() == ['aardvarks', 'zebras']
print('all tests pass; output file is test.dat')
if __name__ == '__main__':
test()