diff options
Diffstat (limited to 'mp_tool')
| -rw-r--r-- | mp_tool/__init__.py | 11 | ||||
| -rw-r--r-- | mp_tool/util.py | 1 | ||||
| -rw-r--r-- | mp_tool/web.py | 116 |
3 files changed, 111 insertions, 17 deletions
diff --git a/mp_tool/__init__.py b/mp_tool/__init__.py index 77b8820..63fe14f 100644 --- a/mp_tool/__init__.py +++ b/mp_tool/__init__.py @@ -1,9 +1,12 @@ - class Constants: ENTER_RAW_MODE = b'\x01' # CTRL-A - ENTER_REPL_MODE = b'\x02' # CTRL-B - INTERRUPT = b'\x03' # CTRL-C - CTRL_D = b'\x04' # CTRL-D + ENTER_REPL_MODE = b'\x02' # CTRL-B + INTERRUPT = b'\x03' # CTRL-C + CTRL_D = b'\x04' # CTRL-D MARKER_BEGIN = b'>>>>>>>>>>' MARKER_END = b'<<<<<<<<<<' + WEBREPL_REQ_S = "<2sBBQLH64s" + WEBREPL_PUT_FILE = 1 + WEBREPL_GET_FILE = 2 + WEBREPL_GET_VER = 3 diff --git a/mp_tool/util.py b/mp_tool/util.py index 7e3f610..b289124 100644 --- a/mp_tool/util.py +++ b/mp_tool/util.py @@ -1,6 +1,7 @@ import argparse import platform + class HelpAction(argparse._HelpAction): def __call__(self, parser, namespace, values, option_string=None): formatter = parser._get_formatter() diff --git a/mp_tool/web.py b/mp_tool/web.py index 6dcccfb..39e3010 100644 --- a/mp_tool/web.py +++ b/mp_tool/web.py @@ -10,23 +10,21 @@ import termios from threading import Thread from sys import stdout, stdin from copy import copy +import os +import struct -def get(url: str, password: str): - raise NotImplementedError() - - -def put(url: str, password: str): - raise NotImplementedError() - - -def connect_and_auth(url, password) -> websocket.WebSocket: +def _connect_and_auth(url: str, password: str) -> websocket.WebSocket: ws = websocket.create_connection(url, timeout=0.5) frame = ws.recv_frame() + if password is None: + stdout.write('Password: ') + stdout.flush() + password = stdin.readline() + if frame.data != b"Password: ": raise Exception("Unexpected response: {}".format(frame.data)) - stdout.write(frame.data.decode('utf-8')) ws.send(password + "\n") frame = ws.recv_frame() @@ -35,8 +33,100 @@ def connect_and_auth(url, password) -> websocket.WebSocket: return ws +def get(url: str, password: str, remote_filename: str, target: str): + if target: + if os.path.isdir(target): + local_filename = os.path.join(target, os.path.basename(remote_filename)) + else: + local_filename = target + else: + local_filename = os.path.basename(remote_filename) + + remote_filename_b = remote_filename.encode('utf-8') + + ws = _connect_and_auth(url, password) + ws.settimeout(5) + + rec = struct.pack(Constants.WEBREPL_REQ_S, b"WA", + Constants.WEBREPL_GET_FILE, 0, 0, 0, + len(remote_filename_b), remote_filename_b) + + ws.send(rec, websocket.ABNF.OPCODE_BINARY) + + frame = ws.recv_frame() + sig, code = struct.unpack("<2sH", frame.data) + if frame.opcode != websocket.ABNF.OPCODE_BINARY or sig != b'WB' or code != 0: + raise Exception("Error initial response sig={} code={}".format(sig, code)) + + ret = b"" + while True: + # Confirm message + ws.send(b"\1", websocket.ABNF.OPCODE_BINARY) + frame = ws.recv_frame() + print(frame) + (sz,) = struct.unpack("<H", frame.data[:2]) + if sz == 0: + break + ret += frame.data[2:] + + frame = ws.recv_frame() + sig, code = struct.unpack("<2sH", frame.data) + if frame.opcode != websocket.ABNF.OPCODE_BINARY or sig != b'WB' or code != 0: + raise Exception("Error final response sig={} code={}".format(sig, code)) + + with open(local_filename, 'wb') as fh: + fh.write(ret) + + +def put(url: str, password: str, local_filename: str, target: str): + if target: + remote_filename = os.path.join(target, local_filename) + else: + remote_filename = os.path.basename(local_filename) + remote_filename_b = remote_filename.encode('utf-8') + + with open(local_filename, 'br') as file_fh: + data = file_fh.read() + + ws = _connect_and_auth(url, password) + ws.settimeout(5) + + sz = len(data) + rec = struct.pack(Constants.WEBREPL_REQ_S, b"WA", + Constants.WEBREPL_PUT_FILE, 0, 0, sz, + len(remote_filename_b), remote_filename_b) + + ws.send(rec[:10], websocket.ABNF.OPCODE_BINARY) + ws.send(rec[10:], websocket.ABNF.OPCODE_BINARY) + + frame = ws.recv_frame() + sig, code = struct.unpack("<2sH", frame.data) + if frame.opcode != websocket.ABNF.OPCODE_BINARY or sig != b'WB' or code != 0: + raise Exception("Error initial response sig={} code={}".format(sig, code)) + + i = 0 + n = int(sz / 256) + 1 + print("{}..{}: ".format(i, n), end="") + cnt = 0 + while True: + buf = data[cnt:cnt + 256] + if not buf: + break + ws.send(buf, websocket.ABNF.OPCODE_BINARY) + + cnt += len(buf) + i += 1 + print("{} ".format(i), end="") + print(". done") + + frame = ws.recv_frame() + sig, code = struct.unpack("<2sH", frame.data) + if frame.opcode != websocket.ABNF.OPCODE_BINARY or sig != b'WB' or code != 0: + raise Exception("Error initial response sig={} code={}".format(sig, code)) + + def eval(url: str, password: str, code: str): - ws = connect_and_auth(url, password) + ws = _connect_and_auth(url, password) ws.send(Constants.ENTER_REPL_MODE) stdout.write(read_until_eval_or_timeout(ws)) ws.send(code + "\r\n") @@ -87,9 +177,9 @@ def set_tty_mode(fd, mode): termios.tcsetattr(fd, termios.TCSAFLUSH, mode) -def repl(url: str, password: str): +def repl(url: str, password: str = None): print("Type ^[ CTRL-] or CTRL-D to quit") - ws = connect_and_auth(url, password) + ws = _connect_and_auth(url, password) ws.send("\x02") reader = Reader(ws) |