#!/usr/bin/env python3

import os
import socket
import asyncio
from string import ascii_letters, digits

FLAG = os.environ.get("FLAG", "olmath{flag_is_not_here}").encode()

REQUEST_TOKEN  = b"\xFF\xAA"
RESPONCE_TOKEN = b"\xFF\xBB"

REQUEST_CODE_LOGIN  = 0x01
REQUEST_CODE_FLAG   = 0x02
REQUEST_CODE_DEBUG  = 0x03

RESPONCE_CODE_TOKEN = 0x64
RESPONCE_CODE_FLAG  = 0xc8
RESPONCE_CODE_ERROR = 0xff

ERROR_CODE_NOT_REQUEST  = 0x01
ERROR_CODE_BAD_PASSWORD = 0x02
ERROR_CODE_BAD_TOKEN    = 0x03
ERROR_CODE_BAD_CODE     = 0x04

__tokens__ = []

def validate_pass(passwd: bytes) -> bool:
    len_pass = len(passwd)
    if (len_pass < 20):
        return False

    s = set(passwd)
    if (len(s) < 6):
        return False

    printable_chars = set(bytes((ascii_letters+digits).encode()))
    if not all(c in printable_chars for c in passwd):
        return False

    sum = 0
    for i in passwd:
        sum ^= ( i % 256 )

    if (sum != 0x56):
        return False

    return True

def generate_token(log, passw):
    token = b''
    for c in range(0,10):
        token += (passw[c] ^ log[c % len(log)] | 0xf0).to_bytes(1, "big")
    return token

def parse_marker(package: bytes) -> bytes:
    return package[0:2]

def parse_code(package: bytes) -> int:
    return int.from_bytes(package[2:3], 'big')

def parse_package_len(package: bytes) -> int:
    return int.from_bytes(package[3:5], 'big')

def parse_login_len(package: bytes) -> int:
    return int.from_bytes(package[5:7], 'big')

def parse_password_len(package: bytes) -> int:
    return int.from_bytes(package[7:9], 'big')

def parse_string(package: bytes, addr: int, length: int) -> bytes:
    return package[addr:addr+length]

def generate_responce(code: int, data: bytes) -> bytes:
    return RESPONCE_TOKEN + code.to_bytes(1, "big") + data

def generate_error_responce(code: int) -> bytes:
    return generate_responce(RESPONCE_CODE_ERROR, code.to_bytes(1, "big"))

async def handle_client(reader, writer):
    data = await reader.read(128000)
    
    code = parse_code(data)
    if (parse_marker(data) != REQUEST_TOKEN):
        writer.write(generate_error_responce(ERROR_CODE_NOT_REQUEST))
        writer.close()
        return

    if code == REQUEST_CODE_LOGIN:
        login_len  = parse_login_len(data)
        passwd_len = parse_password_len(data)
        login      = parse_string(data, 0x09, parse_login_len(data))
        passwd     = parse_string(data, 0x09+parse_login_len(data), parse_password_len(data))
        
        if validate_pass(passwd):
            token = generate_token(login, passwd)
            if not token in __tokens__:
                __tokens__.append(token)
            writer.write(generate_responce(RESPONCE_CODE_TOKEN, token))
        else:
            writer.write(generate_error_responce(ERROR_CODE_BAD_PASSWORD))

    elif code == REQUEST_CODE_FLAG:
        user_token = parse_string(data, 0x05, 10)
        if user_token in __tokens__:
            writer.write(generate_responce(RESPONCE_CODE_FLAG, FLAG))
            __tokens__.remove(user_token)
        else:
            writer.write(generate_error_responce(ERROR_CODE_BAD_TOKEN))

    elif code == REQUEST_CODE_DEBUG:
        str1 = "Request Marker: {}\n".format(str(parse_marker(data)))
        str2 = "Request Debug code: {}\n".format(str(parse_code(data)))
        str3 = "Package Length: {}\n".format(str(parse_package_len(data)))
        data_len = int.from_bytes(parse_string(data, 0x05, 2), "big")
        str4 = "Data Length: {}\n".format(data_len)
        str5 = "Data: {}\n".format(str(parse_string(data, 0x07, data_len)))
        writer.write((str1+str2+str3+str4+str5).encode())

    else:
        writer.write(generate_error_responce(ERROR_CODE_BAD_CODE))

    await writer.drain()
    writer.close()


##################################################################
# Тут идет немного асинхронной магии, к таску не имеет отношения #
##################################################################

async def init():

    sock = await asyncio.start_server(handle_client, '0.0.0.0', 1337)

    async with sock:
        await sock.serve_forever()

if __name__ == "__main__":
    asyncio.run(init())