From 798e0c34aa06c4f4a86dae0064c3558f554ccc78 Mon Sep 17 00:00:00 2001 From: Samuel Sloniker Date: Fri, 16 Jul 2021 16:35:55 -0700 Subject: [PATCH] First step of server redesign Only thing left is to re-add disconnect handling --- server/crops.json | 1 - server/imgproc.py | 40 +------ server/wss.py | 265 +++++++++++++++++----------------------------- 3 files changed, 102 insertions(+), 204 deletions(-) delete mode 100644 server/crops.json diff --git a/server/crops.json b/server/crops.json deleted file mode 100644 index 815c5d9..0000000 --- a/server/crops.json +++ /dev/null @@ -1 +0,0 @@ -["232x53+0+0", "232x11+0+53", "232x84+0+64", "163x148+232+0", "181x148+395+0", "170x148+576+0", "54x148+746+0", "139x148+0+148", "139x184+0+296", "110x166+139+148", "110x166+249+148", "110x166+359+148", "110x166+469+148", "110x166+579+148", "111x166+689+148", "110x166+139+314", "110x166+249+314", "110x166+359+314", "110x166+469+314", "110x166+579+314", "111x166+689+314"] diff --git a/server/imgproc.py b/server/imgproc.py index 9ba9ddf..ab93021 100644 --- a/server/imgproc.py +++ b/server/imgproc.py @@ -3,48 +3,14 @@ import threading import os import subprocess -def run(bmp_path, crop): - subprocess.run(['convert', bmp_path, '-crop', crop, f'pieces/{crop.split("+", 1)[1].replace("+", "x")}.jpg',]) - -def get_split_imgs(): - bmp_path = backend.get_img() - - with open('crops.json') as f: - crops = json.load(f) - - threads = [] - for crop in crops: - threads.append(threading.Thread(target=run, args=([bmp_path, crop]))) - threads[-1].start() - for thread in threads: - thread.join() - - os.unlink(bmp_path) - - with open('newlist', 'w+b') as f: - md5sum = subprocess.Popen(['md5sum', '-c', 'oldlist'], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) - grep = subprocess.Popen(['grep', 'FAILED'], stdin=md5sum.stdout, stdout=f) - grep.wait() - - changed = [] - with open('newlist') as f: - lines = f.readlines() - for line in lines: - changed.append(line.split(": ")[0].split('.')[0].split('/')[1]) - - with open('oldlist', 'w+b') as f: - subprocess.Popen(['md5sum'] + [f'pieces/{i}' for i in os.listdir('pieces')], stdout=f) - - return changed - -def get_full_img(): +def get_img(): bmp_path = backend.get_img() - subprocess.run(['convert', bmp_path, 'img.jpg',]) + subprocess.run(['convert', bmp_path, '-define', 'webp:lossless=true', 'img.webp',]) os.unlink(bmp_path) - return 'img.jpg' + return 'img.webp' def touch(x, y, w, is_long): x = round(800 * x / w) diff --git a/server/wss.py b/server/wss.py index 53d2650..c32c7be 100644 --- a/server/wss.py +++ b/server/wss.py @@ -1,19 +1,14 @@ import argon2 -import asyncio import base64 import imgproc import importlib +import logging import os import parse_config -import random -import requests -import secrets -import shutil import sys -import tempfile import threading import time -import tornado.web, tornado.websocket, tornado.ioloop +from websocket_server import WebsocketServer try: conf = sys.argv[1] @@ -23,198 +18,136 @@ except IndexError: with open(conf) as f: config_data = parse_config.load(f) -imgproc.backend = importlib.import_module(f'backends.{config_data["backend"]}') -imgproc.backend.config = config_data +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 -token = secrets.token_urlsafe(64) - - class DisconnectError(BaseException): pass class Client: - def __init__(self, conn): + def __init__(self, server, client_): global client if client is not None: raise DisconnectError('err%*inuse%Server already in use') client = self - self.conn = conn - self.items = {} + self.server = server + self.client = client_ + #self.conn = conn self.lock = threading.Lock() - self.good = True - - def send(self, item, name): - with self.lock: - self.items[name] = item - - def ack(self): - with self.lock: - self.good = False - - def get_item_to_send(self): - with self.lock: - if self.items: - name, content = list(self.items.items())[0] - del self.items[name] - return name, content - elif not self.good: - return 'ack', 'ack' - else: - return None, None - - - -class HCRAServer(tornado.websocket.WebSocketHandler): - def open(self): self.has_auth = False + self.ready_for_msgs = False + print('srfm') + self.acked = True self.version = None + self.next_msg = None - def on_close(self): + def close(self): global client - if client is self.client: - self.client.good = None - client = None - self.is_open = False + self.ready_for_msgs = False + client = None imgproc.backend.disconnect() - def on_message(self, message): - action = message.split(' ', 1)[0] - if self.version is None: - if action == 'maxver': - self.version = min(int(message.split(' ', 1)[1]), 1) - self.write_message(f'ver%{self.version}') - else: - self.write_message('err%*mustmaxver%Client must send version') - self.close() - return - elif not self.has_auth: - if action != 'pass': - self.write_message('err%*mustauth%Authentication required') - self.close() - return + 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 - try: - ph.verify(config_data['password_argon2'], message.split(' ', 1)[1]) - except argon2.exceptions.VerifyMismatchError: - self.write_message(f'err%*badpass%Incorrect password') - self.close() - return + def got_ack(self): + self.acked = True + self.do_send() - self.has_auth = True - try: - self.client = Client(self) - except DisconnectError as e: - self.write_message(str(e)) - return - - imgproc.backend.connect() - - try: - imgname = imgproc.get_full_img() - except Exception as e: - self.write_message('err%noconn%Server failed to capture screenshot') - return - - with open(imgname, 'rb') as f: - img = f.read() - os.unlink(imgname) - self.write_message(f'pic%0x0%data:image/jpeg;base64,{base64.b64encode(img).decode("utf-8")}') - self.is_open = True - self.client.ack() - - else: - if action == 'ack': - self.client.good = True - 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 check_origin(self, origin): - return True - -def cycle(): - try: - changed = imgproc.get_split_imgs() - except Exception as e: - if client is not None: - client.send('err%noconn%Server failed to capture screenshot', 'ERR') - time.sleep(3) +def on_message(client_, server, message): + global client + if client_ != client.client: return - threads = [] - for i in changed: - thread = threading.Thread(target=do_img, args=(i,)) - threads.append(thread) - thread.start() + 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: - client.ack() - - -def do_img(imgname): - if client is not None: - client.send(img(imgname), imgname) - - -def img(imgname): - with open(f'pieces/{imgname}.jpg', 'rb') as f: - img = f.read() - - response = f'pic%{imgname}%data:image/jpeg;base64,{base64.b64encode(img).decode("utf-8")}' - - return response - + msg = get_img_msg() + client.next_msg = msg + client.do_send() def do_cycles(): while True: - if client is not None: - cycle() - time.sleep(0.4) - else: - time.sleep(1) + cycle() + time.sleep(0.8) - - -class Cycler(tornado.web.RequestHandler): - def get(self): - if client is not None: - name, item = client.get_item_to_send() - if item is not None: - client.conn.write_message(item) - self.write('OK') - - -def get_token(): - while True: - requests.get('http://localhost:1234/' + token) - - -tmp = tempfile.mkdtemp(prefix="HCRA-") -try: - shutil.copy('crops.json', tmp) - os.chdir(tmp) - os.mkdir('pieces') - - threading.Thread(target=do_cycles).start() - - application = tornado.web.Application([ - (r"/", HCRAServer), - ('/' + token, Cycler), - ]) - application.listen(int(config['port'])) - threading.Thread(target=get_token).start() - tornado.ioloop.IOLoop.current().start() -finally: - os.chdir('/') - shutil.rmtree(tmp) +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.run_forever()