hcra/server/ws_server/wss.py

import argon2
import base64
import imgproc
import importlib
import logging
import os
import parse_config
import sys
import threading
import time
from websocket_server import WebsocketServer

try:
    conf = sys.argv[1]
except IndexError:
    conf = 'conf.txt'

with open(conf) as f:
    config_data = parse_config.load(f)

backend = importlib.import_module(f'backends.{config_data["backend"]}')
backend.config = config_data
imgproc.config = config_data
imgproc.backend = backend

ph = argon2.PasswordHasher()

client = None


class DisconnectError(BaseException):
    pass


class Client:
    def __init__(self, server, client_):
        global client
        if client is not None:
            raise DisconnectError('err%*inuse%Server already in use')

        client = self
        self.server = server
        self.client = client_
        self.has_auth = False
        self.ready_for_msgs = False
        self.acked = True
        self.version = None
        self.next_msg = None

    def close(self):
        global client
        self.ready_for_msgs = False
        client = None
        imgproc.backend.disconnect()

    def do_send(self):
        if self.next_msg is not None and self.ready_for_msgs and self.acked:
            msg = self.next_msg
            self.next_msg = None
            self.server.send_message(self.client, msg)
            self.server.send_message(self.client, 'ack')
            self.acked = False

    def got_ack(self):
            self.acked = True
            self.do_send()


def on_left(client_, server):
    global client
    if client_ != client.client:
        return
    
    client.close()


def on_message(client_, server, message):
    global client
    if client_ != client.client:
        return

    action = message.split(' ', 1)[0]
    if client.version is None:
        if action == 'maxver':
            client.version = min(int(message.split(' ', 1)[1]), 1)
            server.send_message(client_, f'ver%{client.version}')
        else:
            server.send_message(client_, 'err%*mustmaxver%Client must send version')
            client_['handler'].send_text("", opcode=0x8)
            return
    elif not client.has_auth:
        if action != 'pass':
            server.send_message(client_, 'err%*mustauth%Authentication required')
            client_['handler'].send_text("", opcode=0x8)
            return

        try:
            ph.verify(config_data['password_argon2'], message.split(' ', 1)[1])
        except argon2.exceptions.VerifyMismatchError:
            server.send_message(client_, f'err%*badpass%Incorrect password')
            client_['handler'].send_text("", opcode=0x8)
            return

        client.has_auth = True
        
        imgproc.backend.connect()

        client.ready_for_msgs = True
        client.do_send()

    else:
        if action == 'ack':
            client.got_ack()
        else:
            _, x, y, w, is_long = message.split(' ')
            x, y, w, is_long = int(x), int(y), int(w), is_long == 'true'
            imgproc.touch(x, y, w, is_long)


def on_connect(client_, server):
    try:
        client = Client(server, client_)
    except DisconnectError as e:
        server.send_message(client_, str(e))
        client_['handler'].send_text("", opcode=0x8)
        return


def get_img_msg():
    try:
        img = imgproc.get_img()
    except Exception as e:
        return 'err%noconn%Server failed to capture screenshot'

    with open(img, 'rb') as f:
        img_data = f.read()
    
    os.unlink(img)

    return f'pic%0x0%data:image/webp;base64,{base64.b64encode(img_data).decode("utf-8")}'

def cycle():
    if client is not None:
        msg = get_img_msg()
        try:
            client.next_msg = msg
            client.do_send()
        except AttributeError as e:
            if "'NoneType' object has no attribute" in str(e):
                return
            else:
                raise e

def do_cycles():
    while True:
        cycle()
        time.sleep(0.8)


threading.Thread(target=do_cycles).start()
server = WebsocketServer(int(config_data['port']), host='0.0.0.0', loglevel=logging.INFO)
server.set_fn_new_client(on_connect)
server.set_fn_message_received(on_message)
server.set_fn_client_left(on_left)
server.run_forever()