196 lines
4.6 KiB
Python
196 lines
4.6 KiB
Python
import argon2
|
|
import asyncio
|
|
import base64
|
|
import imgproc
|
|
import importlib
|
|
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
|
|
|
|
try:
|
|
conf = sys.argv[1]
|
|
except IndexError:
|
|
conf = 'conf.txt'
|
|
|
|
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
|
|
imgproc.config = config_data
|
|
|
|
ph = argon2.PasswordHasher()
|
|
|
|
client = None
|
|
|
|
|
|
token = secrets.token_urlsafe(64)
|
|
|
|
|
|
class DisconnectError(BaseException):
|
|
pass
|
|
|
|
|
|
class Client:
|
|
def __init__(self, conn):
|
|
global client
|
|
if client is not None:
|
|
raise DisconnectError('err%*inuse%Server already in use')
|
|
|
|
client = self
|
|
self.conn = conn
|
|
self.items = {}
|
|
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):
|
|
try:
|
|
self.client = Client(self)
|
|
except DisconnectError as e:
|
|
self.write_message(str(e))
|
|
return
|
|
|
|
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()
|
|
|
|
def on_close(self):
|
|
global client
|
|
if client is self.client:
|
|
self.client.good = None
|
|
client = None
|
|
self.is_open = False
|
|
|
|
def on_message(self, message):
|
|
action = message.split(' ', 1)[0]
|
|
if action == 'ack':
|
|
self.client.good = True
|
|
else:
|
|
_, password, x, y, w, is_long = message.split(' ')
|
|
try:
|
|
ph.verify(config_data['password_argon2'], password)
|
|
x, y, w, is_long = int(x), int(y), int(w), is_long == 'true'
|
|
imgproc.touch(x, y, w, is_long)
|
|
except argon2.exceptions.VerifyMismatchError:
|
|
self.client.send(f'err%*badpass%Incorrect password', 'BADPASS')
|
|
self.close()
|
|
|
|
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)
|
|
return
|
|
|
|
threads = []
|
|
for i in changed:
|
|
thread = threading.Thread(target=do_img, args=(i,))
|
|
threads.append(thread)
|
|
thread.start()
|
|
|
|
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
|
|
|
|
|
|
def do_cycles():
|
|
while True:
|
|
if client is not None:
|
|
cycle()
|
|
time.sleep(0.4)
|
|
else:
|
|
time.sleep(1)
|
|
|
|
|
|
|
|
|
|
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(1234)
|
|
threading.Thread(target=get_token).start()
|
|
tornado.ioloop.IOLoop.current().start()
|
|
finally:
|
|
os.chdir('/')
|
|
shutil.rmtree(tmp)
|