orbital/tools/dumper/server.py
Alexander Richards 78f1cc81e2
Update server.py
2019-06-24 16:53:57 +02:00

121 lines
3.6 KiB
Python

#!/usr/bin/env python3
import argparse
import binascii
import socket
import struct
import threading
import os
import sys
import aiohttp
from aiohttp import web
# Configuration
PORT_BLOBS = 9021
PORT_DEBUG = 9022
# Context
current_file = None
# Sockets
def server_blobs():
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
s.bind(('', PORT_BLOBS))
except socket.error as msg:
print('blobs-server: Bind failed: %s\n' % msg)
sys.exit()
s.listen(5)
while True:
c, addr = s.accept()
print('blobs-server: Client connected: %s:%s' % addr)
while True:
# File path
path_size = c.recv(8, socket.MSG_WAITALL)
if not path_size: break
path_size = struct.unpack('Q', path_size)[0]
if not path_size: break
path = c.recv(path_size, socket.MSG_WAITALL)
path = os.path.join('dump', path.decode('utf-8'))
# File data
data_size = c.recv(8, socket.MSG_WAITALL)
if not data_size: break
data_size = struct.unpack('Q', data_size)[0]
if not data_size: break
data = c.recv(data_size, socket.MSG_WAITALL)
# Save file
path_dir = os.path.dirname(path)
if path_dir and not os.path.exists(path_dir):
os.makedirs(path_dir, exist_ok=True)
with open(path, 'wb') as f:
f.write(data)
print('blobs-server: Client disconnected: %s:%s' % addr)
c.close()
s.close()
def server_debug():
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
s.bind(('', PORT_DEBUG))
except socket.error as msg:
print('debug-server: Bind failed: %s\n' % msg)
sys.exit()
s.listen(5)
while True:
c, addr = s.accept()
print('debug-server: Client connected: %s:%s' % addr)
while True:
# TODO: There's surely a better way, but whatever
byte = c.recv(1)
if not byte:
break
sys.stdout.buffer.write(byte)
sys.stdout.flush()
print('debug-server: Client disconnected: %s:%s' % addr)
c.close()
s.close()
# Website
async def handle_index(request):
return web.FileResponse('./index.html')
async def handle_websocket(request):
ws = web.WebSocketResponse()
await ws.prepare(request)
async for msg in ws:
if msg.type == aiohttp.WSMsgType.TEXT:
current_file = os.path.normpath(
os.path.join('dump', msg.data))
elif msg.type == aiohttp.WSMsgType.BINARY and current_file:
os.makedirs(os.path.dirname(current_file), exist_ok=True)
with open(current_file, 'wb') as f:
f.write(msg.data)
elif msg.type == aiohttp.WSMsgType.ERROR:
print('WS connection closed with exception %s' % ws.exception())
print('WS connection closed')
return ws
def main():
# Handle arguments
parser = argparse.ArgumentParser(
description='Create server for Orbital dumper.')
parser.add_argument('-p', '--port', type=int, default=80, required=False,
help='Port for HTTP/WS server')
args = parser.parse_args()
# Create sockets
t_blobs = threading.Thread(target=server_blobs)
t_debug = threading.Thread(target=server_debug)
t_blobs.start()
t_debug.start()
# Create webserver
app = web.Application()
app.router.add_get('/', handle_index)
app.router.add_get('/ws', handle_websocket)
app.router.add_static('/', path='.')
web.run_app(app, port=args.port)
if __name__ == '__main__':
main()