main.py

#! /usr/bin/env python3

from collections import namedtuple
from websockets import serve

import jwt

import json
import asyncio
import random
import uuid
import re
import os

# FIXME: should support _FILE for secrets
JWT_SECRET = os.environ['ALLOC_JWT_SECRET']
STAFF_PASSWORD = os.environ['ALLOC_STAFF_PW']

GUEST = 0
STUDENT = 1
STAFF = 2

def allocateVariableGroups(students, idealTeamSize=4, minTeamSize=2):
    """Allocate students to groups, with an ideal team size.
    If there are students left over, and there are at least minStudents they get
    to be a team. Else, the 'remainder' students are distrubted across teams.
    """
    random.shuffle( students )
    groups = []

    # Allocate teams to groups
    gid = 0
    currGroup = []
    for student in students:
        currGroup.append( student )
        if idealTeamSize == len(currGroup):
            groups.append( currGroup )
            currGroup = []
            gid += 1

    # last team
    if currGroup:
        if minTeamSize <= len(currGroup) or len(groups) == 0:
            # viable team (or no other choice)
            groups.append( currGroup )
        else:
            # redistribute across teams
            groupAlloc = 0
            while currGroup:
                groups[groupAlloc].append( currGroup.pop() )
                groupAlloc = (groupAlloc + 1) % len(groups)

    return groups

def allocateFixedGroups(students, numGroups, minSize = 1):
    """Allocate students assuming a fixed number of groups.

    Students are allocated fairly amoung groups, remainders
    will be added to the first n groups. If there are not
    enouph students for numGroups, then the maximum number
    of groups to still satify numGroups is used.
    """

    # figure out how many groups we can have
    maxPossibleGroups = len(students) // minSize
    numGroups = min( numGroups, maxPossibleGroups )

    if numGroups == 0:
        numGroups = 1 # we're having a bad day

    groups = []
    for i in range(0, numGroups):
        groups.append( [] )

    random.shuffle( students )
    nextGroup = 0
    while students:
        groups[nextGroup].append( students.pop() )
        nextGroup = (nextGroup + 1) % numGroups

    return groups

ALLOCATIONS = {
    'members': allocateVariableGroups,
    'group': allocateFixedGroups
}

async def sendToStaff( topic, data ):
    for staff in get_staff():
        await clients[staff].send_topic( topic, data )

async def changeName(client, message):
    realName = message['data']['name']
    realName = re.sub(r'[^\w -]', '', realName)

    if 'password' in message['data']:
        if message['data']['password'] == STAFF_PASSWORD:
            client.role = STAFF
        else:
            client.role = GUEST
    else:
        client.role = STUDENT

    await client.rename( realName )

async def forceAllocate(client, message):
    if client.role != STAFF:
        return

    data = message["data"]

    # Allow different group sizes
    size = 4
    if 'idealSize' in data:
        size = int(data['idealSize'])

    minSize = 2
    if 'minSize' in data:
        minSize = int(data['minSize'])
    
    method = allocateFixedGroups
    if 'method' in data and data['method'] in ALLOCATIONS:
        method = ALLOCATIONS[ data['method'] ]

    summary = {
        'size': size,
        'groups': []
    }

    # Perform group allocation
    ids = [ x for x in clients.keys() if clients[x].role == STUDENT ]
    groups = method( ids, size, minSize )

    # send out group notifications
    for (idx, group) in enumerate(groups):
        members = [ str( clients[x].name ) for x in group ]
        data = {
            'groupID': idx,
            'members': members
        }
        summary['groups'].append( data )
        for student in group:
            await clients[ student ].send_topic( 'allocate', data )

    # Send master list to staff
    await sendToStaff( 'groups', summary )

async def resumeToken(client, message):
    try:
        ticket = jwt.decode( message['data'], JWT_SECRET, algorithms=["HS256"] )
        resumed_token = uuid.UUID("{"+ticket['uuid']+"}")
        del clients[ client.uuid ]
        print( "client {} used jwt to reclaim id {}", client.uuid, resumed_token )

        # resume session
        client.uuid = resumed_token
        clients[ client.uuid ] = client
        if resumed_token in names:
            session = names[ resumed_token ]
            client.name = session[0]
            client.role = session[1]
            await client.send_topic('changeName', { 'name': client.name, 'role': client.role })

            if client.role == STAFF:
                students = [ {'uuid': str(x.uuid), 'name': x.name} for x in clients.values() if x.role == STUDENT ]
                await client.send_topic('students', students)
    except jwt.exceptions.DecodeError:
        pass

async def forceRename(client, message):
    if client.role != STAFF:
        return

    realName = message['data']['name']
    realName = re.sub(r'[^\w -]', '', realName)
    client_token = uuid.UUID("{"+message['data']['uuid']+"}")

    await clients[ client_token ].rename( realName )

protocol = {
        'changeName': changeName,
        'token': resumeToken,

        # staff options
        'forceAllocate': forceAllocate,
        'forceRename': forceRename,
}

class Client:

    def __init__(self, socket):
        self.uuid = uuid.uuid4()
        self.name = None
        self.role = GUEST
        self.socket = socket

    async def rename(self, new_name):
        self.name = new_name
        names[ self.uuid ] = ( new_name, self.role )
        await self.send_topic('changeName', {'name': new_name, 'role': self.role })

    async def send_topic(self, topic, data):
        msg = {
            'type': topic,
            'data': data
        }
        print( "[>>]", msg )
        await self.socket.send( json.dumps(msg) )

    async def recv_topic(self, data):
        print("[<<]", data)
        msg_type = data['type']
        if msg_type in protocol:
            await protocol[ msg_type ]( self, data )

clients = {}
names = {}

def get_students():
    return [ x for x in clients.keys() if clients[x].role == STUDENT ]

def get_staff():
    return [ x for x in clients.keys() if clients[x].role == STAFF ]

async def handleClient(websocket):
    client = Client(websocket)
    clients[ client.uuid ] = client

    auth_token = jwt.encode( {'uuid': str(client.uuid)}, JWT_SECRET, algorithm="HS256" )
    await client.send_topic("token", auth_token)

    async for message in websocket:
        data = json.loads( message )
        msg_type = data['type']
        if msg_type in protocol:
            await client.recv_topic( data )

    del clients[ client.uuid ]

async def main():
    async with serve(handleClient, "0.0.0.0", 8081):
        await asyncio.Future()

asyncio.run( main() )