user_wrapper

#! /usr/bin/env python3

import os
import sys
import re
import shutil

from collections import defaultdict
from pathlib import Path
from hashlib import sha256

##### CONSTANTS #####

SECRET = 'th1s_s3cr3t_pr0t3cts_th3_fl4gs'

ALIAS_RE = '^[a-z0-9_]{2,15}$'
ASURITE_RE = '^[a-z]{2,15}[0-9]{0,3}$'

ENROLLED = {

}

##### ASSIGNMENTS #####

class ChallengeSet:
    def __init__(self, docker_name, binary_path_re, score_multiplier=1, docker_arguments=None, custom_challenges=True):
        if not docker_arguments:
            docker_arguments = []

        self.docker_name = docker_name
        self.binary_path_re = binary_path_re
        self.score_multiplier = score_multiplier
        self.docker_arguments = docker_arguments
        self.custom_challenges = custom_challenges

    def load(self):
        global DOCKER_NAME
        global BINARY_PATH_RE
        global SCORE_MULTIPLIER
        global DOCKER_ARGUMENTS
        global CUSTOM_CHALLENGES

        DOCKER_NAME = self.docker_name
        BINARY_PATH_RE = self.binary_path_re
        SCORE_MULTIPLIER = self.score_multiplier
        DOCKER_ARGUMENTS = self.docker_arguments
        CUSTOM_CHALLENGES = self.custom_challenges

        
HOMEWORKS = {
    'hw1': ChallengeSet('hw1', '^/[a-zA-Z0-9/_\-\+\.]+$', 1, custom_challenges=False),
    'hw2': ChallengeSet('hw2', '^/[a-zA-Z0-9/_\-\+\.]+$', 2, custom_challenges=False),
    'hw3': ChallengeSet('hw3', '^/pwn/[a-z0-9_]+$', 3),
    'hw4': ChallengeSet('hw4', '^/pwn/[a-z0-9_]+$', 4),
    'hw5': ChallengeSet('hw5', '^/pwn/(babypwn|babykey)/[a-z0-9_]+$', 3),
    'hw6': ChallengeSet('hw6', '^/pwn/(babyshell)/[a-z0-9_]+$', 4),
    'hw7': ChallengeSet('hw7', '^/pwn/[a-z0-9_]+$', 2),
}

##### GLOBAL SESSION INFO #####

alias = None
asurite = None

    
##### USER CONTROLS #####

def login():
    alias = input("Hacker Alias: ")
    asurite = input("ASURITE: ")

    if not re.match(ALIAS_RE, alias):
        raise Exception(f"Hacker Alias must match: {ALIAS_RE}")
    if not re.match(ASURITE_RE, asurite):
        raise Exception(f"ASURITE must match: {ASURITE_RE}")

    users = dict()
    for path in Path.home().iterdir():
        if path.is_dir() and ':' in path.name:
            existing_alias, existing_asurite = path.name.split(':', 1)
            users[existing_alias] = existing_asurite

    current_asurite = users.get(alias)
    if current_asurite == asurite:
        return alias, asurite # User exists, successful login
    elif alias not in users.keys() and asurite not in users.values():
        user_path = Path.home() / f'{alias}:{asurite}'
        user_path.mkdir()
        (user_path / 'logs').mkdir()
        (user_path / 'solves').mkdir()

        if CUSTOM_CHALLENGES:
            try:
                challenge = list((Path.home() / 'challenges').iterdir())[0]
                shutil.move(str(challenge), str(user_path / 'challenges'))
            except:
                raise Exception("Error: Yan's fault")
        else:
            (user_path / 'challenges').mkdir()
                

        return alias, asurite # User successfully created
    elif alias in users.keys():
        raise Exception("Hacker Alias already registered!")
    elif asurite in users.values():
        raise Exception("ASURITE already registered!")
    else:
        raise Exception("???")

def show_scoreboard():
    user_scores = {user_asurite[0]: (SCORE_MULTIPLIER * len(solves)) for user_asurite, solves in solves().items()}
    user_grades = grades({user_asurite[0]: (SCORE_MULTIPLIER * len(solves)) for user_asurite, solves in solves().items() if user_asurite[1] in ENROLLED})

    rank_length = max(len(str(len(user_scores))), len('Rank'))
    alias_length = max(*(len(user) for user in user_scores), len('Hacker'))
    score_length = max(*(len(str(user_scores[user])) for user in user_scores), len('Score'))
    grade_length = max(*(len(str(user_grades[user])) for user in user_grades), len('Grade'), len('GUEST'))

    print()
    print("=" * 20 + " SCOREBOARD " + "=" * 20)
    print(f"{'Rank':>{rank_length}}     {'Hacker':<{alias_length}}    {'Score':>{score_length}}    {'Grade':>{grade_length}}")
    for i, alias in enumerate(reversed(sorted(user_scores, key=lambda k: user_scores[k]))):
        grade = user_grades.get(alias, "GUEST")
        print(f"{i+1:>{rank_length}}.    {alias:<{alias_length}}    {user_scores[alias]:>{score_length}}    {grade:>{grade_length}}")
    print()

def solve_challenge():
    challenge = challenge_path()
    base_challenge = os.path.basename(challenge)
    flag = challenge_flag(base_challenge)

    # Kill any old docker container
    filters = ['created', 'paused', 'exited', 'dead']
    cmd = 'docker ps -aq ' + ' '.join(f'--filter "status={f}"' for f in filters) + f' --filter "name={DOCKER_NAME}_{alias}" | xargs -r docker rm'
    os.system(cmd)
    cmd = f'docker ps -q --filter "name={DOCKER_NAME}_{alias}" | xargs -r docker kill'
    os.system(cmd)

    challenges = str(Path.home() / f'{alias}:{asurite}' / 'challenges')

    # Start new docker container
    docker_arguments = [
        ('--name', f'{DOCKER_NAME}_{alias}'),
        ('--env', f'PASSWORD={asurite}'),
        ('--env', f'FLAG={flag}'),
        ('--env', f'BINARY_FILE={challenge}'),
        ('--mount', f'type=bind,source={challenges},destination=/challenges,readonly'),
        ('--mount', f'type=bind,source=/shared,destination=/shared,readonly'),
        ('--publish', '22000-25000:22'),
        ('--cap-add', 'SYS_PTRACE'),
        ('--cpus', '0.5'),
        ('--memory', '500m'),
        ('--memory-swap', '-1'),
        ('--pids-limit', '100'),
        # ('--device-write-bps', '/dev/sda:10kb')
    ]
    docker_arguments += DOCKER_ARGUMENTS
    docker_arguments = ' '.join(f'{param} {value}' for param, value in docker_arguments)

    os.system(f'docker run -td --rm {docker_arguments} {DOCKER_NAME}')
    os.system(f'docker ps --filter="name={DOCKER_NAME}_{alias}"')

    submitted_flag = input("Flag: ")
    if submitted_flag == flag:
        print("Correct Flag!")
        (Path.home() / f'{alias}:{asurite}' / 'solves' / base_challenge).touch()
    else:
        print("Wrong Flag!")

def show_challenges():
    challenges = Path.home() / f'{alias}:{asurite}' / 'challenges'
    for challenge in challenges.glob('**/*'):
        if challenge.is_file():
            print(f"/pwn{str(challenge)[len(str(challenges)):]}")

##### UTILITY FUNCTIONS #####

def challenge_flag(challenge):
    return 'CSE466{' + sha256(f'{SECRET}+{alias}+{challenge}'.encode()).hexdigest() + '}'

def challenge_path():
    challenge = input("Path to Binary: ")
    if not re.match(BINARY_PATH_RE, challenge):
        raise Exception(f"Path to Binary must match: {BINARY_PATH_RE}")
    return challenge

def solves():
    result = dict()
    for path in Path.home().iterdir():
        if path.is_dir() and ':' in path.name:
            current_alias, current_asurite = path.name.split(':', 1)
            result[(current_alias, current_asurite)] = list((path / 'solves').iterdir())
    return result

def grades(user_scores):
    below_grades = {
        user: score
        for user, score in user_scores.items()
        if score <= 70
    }
    above_grades = {
        user: score
        for user, score in user_scores.items()
        if score > 70
    }

    above_score_users = defaultdict(list)
    for user, score in above_grades.items():
        above_score_users[score].append(user)

    clusters = {
        (score, score): above_score_users[score]
        for score in sorted(above_score_users)
    }

    while len(clusters) > 40:
        closest_low = None
        closest_high = None
        closest_distance = None

        prev_cluster = None
        for cluster in sorted(clusters.keys(), key=lambda k: k[0]):
            if closest_low is None:
                closest_low = cluster
            elif closest_high is None:
                closest_high = cluster
                closest_distance = closest_high[0] - closest_low[1]
            else:
                current_distance = cluster[0] - prev_cluster[1]
                if current_distance < closest_distance:
                    closest_low = prev_cluster
                    closest_high = cluster
                    closest_distance = current_distance
            prev_cluster = cluster

        new_cluster = (closest_low[0], closest_high[1])
        new_cluster_data = [*clusters[closest_low], *clusters[closest_high]]

        del clusters[closest_low]
        del clusters[closest_high]
        clusters[new_cluster] = new_cluster_data

    above_grades = dict()
    score = 110
    for cluster in reversed(sorted(clusters.keys(), key=lambda k: k[0])):
        cluster_data = clusters[cluster]
        for user in cluster_data:
            above_grades[user] = score
        score -= 1

    return {**below_grades, **above_grades}

##### MENU #####

def menu():
    global alias
    global asurite

    alias, asurite = login()

    print()
    print()

    if asurite in ENROLLED:
        print("You are enrolled in the course! The scoreboard reflects your grade!")
    else:
        print("You are a guest user! You will not be graded!")

    print()
    print()

    while True:
        print("1. Show Scoreboard")
        print("2. Solve Challenge")
        if CUSTOM_CHALLENGES:
            print("3. Show Challenges")
        print("0. Quit")

        try:
            choice = int(input("Choice: "))
        except ValueError:
            choice = 0

        if choice == 1:
            show_scoreboard()

        elif choice == 2:
            solve_challenge()

        elif CUSTOM_CHALLENGES and choice == 3:
            show_challenges()

        elif choice == 0:
            return

        else:
            print("Invalid Choice!")

        print()
        print()

def main():
    hw = sys.argv[1]
    HOMEWORKS[hw].load()
    menu()

    
if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        print(str(e))