put/get via webrepl

4 files changed, 117 insertions, 19 deletions

diff --git a/mp-tool b/mp-tool
index 1419089..f1823fc 100755
--- a/mp-tool
+++ b/mp-tool
@@ -31,9 +31,11 @@ if __name__ == '__main__':
     parser_repl.add_argument("--password")
 
     parser_put = subparsers.add_parser("put", help="Send file over websocket")
-    parser_put.set_defaults(func=lambda a: import_module('mp_tool.web').put(a.WEBSOCKET, a.password, a.CODE))
+    parser_put.set_defaults(func=lambda a: import_module('mp_tool.web').put(a.WEBSOCKET, a.password, a.FILE, a.TARGET))
     parser_put.add_argument("WEBSOCKET", help="Websocket address (e.g. ws://ESP_E1278E:8266)")
     parser_put.add_argument("--password")
+    parser_put.add_argument("FILE", help="Filename")
+    parser_put.add_argument("TARGET", nargs='?', help="remote target path/filename")
 
     parser_put_serial = subparsers.add_parser("put-serial", help="Send file over serial")
     parser_put_serial.set_defaults(func=lambda a: import_module('mp_tool.serial').put(a.port, a.FILE, a.TARGET))
@@ -42,8 +44,10 @@ if __name__ == '__main__':
     parser_put_serial.add_argument("TARGET", nargs='?', help="remote target path/filename")
 
     parser_get = subparsers.add_parser("get", help="Load file over websocket")
-    parser_get.set_defaults(func=lambda a: import_module('mp_tool.web').get(a.WEBSOCKET, a.password))
+    parser_get.set_defaults(func=lambda a: import_module('mp_tool.web').get(a.WEBSOCKET, a.password, a.FILE, a.TARGET))
     parser_get.add_argument("WEBSOCKET", help="Websocket address (e.g. ws://ESP_E1278E:8266)")
+    parser_get.add_argument("FILE", help="Filename")
+    parser_get.add_argument("TARGET", nargs='?', help="local target path/filename")
     parser_get.add_argument("--password")
 
     parser_get_serial = subparsers.add_parser("get-serial", help="Get file over serial")
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)