Added ability to download files over http(s)

Also fixed some tests and a couple other bugs
This commit is contained in:
Josiah Baldwin
2024-12-12 16:06:18 -08:00
parent 4eda4e6c08
commit a3c721318d
12 changed files with 166 additions and 68 deletions

View File

@@ -4,11 +4,14 @@ from . import exceptions
from . import util
import asyncio
import json
import urllib
import shutil
class Files(tunnel.Tunnel):
def __init__(self, session, nodeid):
super().__init__(session, nodeid, constants.Protocol.FILES)
def __init__(self, session, node):
super().__init__(session, node.nodeid, constants.Protocol.FILES)
self.recorded = None
self._node = node
self._request_id = 0
self._request_queue = asyncio.Queue()
self._download_finished = asyncio.Event()
@@ -16,6 +19,16 @@ class Files(tunnel.Tunnel):
self._current_request = None
self._handle_requests_task = asyncio.create_task(self._handle_requests())
self._chunk_size = 65564
proxies = {}
if self._session._proxy is not None:
# We don't know which protocol the user is going to use, but we only need support one at a time, so just assume both
proxies = {
"http_proxy": self._session._proxy,
"https_proxy": self._session._proxy
}
self._proxy_handler = urllib.request.ProxyHandler(proxies=proxies)
self._http_opener = urllib.request.build_opener(self._proxy_handler, urllib.request.HTTPSHandler(context=self._ssl_context))
def _get_request_id(self):
self._request_id = (self._request_id+1)%(2**32-1)
@@ -68,6 +81,7 @@ class Files(tunnel.Tunnel):
Args:
directory (str): Path to the directory you wish to list
timeout (int): duration in seconds to wait for a response before throwing an error
Returns:
list[~meshctrl.types.FilesLSItem]: The directory listing
@@ -75,6 +89,7 @@ class Files(tunnel.Tunnel):
Raises:
:py:class:`~meshctrl.exceptions.ServerError`: Error from server
:py:class:`~meshctrl.exceptions.SocketError`: Info about socket closure
asyncio.TimeoutError: Command timed out
"""
data = await self._send_command({"action": "ls", "path": directory}, "ls", timeout=timeout)
return data["dir"]
@@ -101,10 +116,12 @@ class Files(tunnel.Tunnel):
Args:
directory (str): Path of directory to create
timeout (int): duration in seconds to wait for a response before throwing an error
Raises:
:py:class:`~meshctrl.exceptions.ServerError`: Error from server
:py:class:`~meshctrl.exceptions.SocketError`: Info about socket closure
asyncio.TimeoutError: Command timed out
Returns:
bool: True if directory was created
@@ -127,10 +144,12 @@ class Files(tunnel.Tunnel):
path (str): Directory from which to delete files
files (str|list[str]): File or files to remove from the directory
recursive (bool): Whether to delete the files recursively
timeout (int): duration in seconds to wait for a response before throwing an error
Raises:
:py:class:`~meshctrl.exceptions.ServerError`: Error from server
:py:class:`~meshctrl.exceptions.SocketError`: Info about socket closure
asyncio.TimeoutError: Command timed out
Returns:
str: Info about the files removed. Something along the lines of Delete: "/path/to/file", or 'Delete recursive: "/path/to/dir", n element(s) removed'.
@@ -155,10 +174,12 @@ class Files(tunnel.Tunnel):
path (str): Directory from which to rename the file
name (str): File to rename
new_name (str): New name to give the file
timeout (int): duration in seconds to wait for a response before throwing an error
Raises:
:py:class:`~meshctrl.exceptions.ServerError`: Error from server
:py:class:`~meshctrl.exceptions.SocketError`: Info about socket closure
asyncio.TimeoutError: Command timed out
Returns:
str: Info about file renamed. Something along the lines of 'Rename: "/path/to/file" to "newfile"'.
@@ -173,6 +194,7 @@ class Files(tunnel.Tunnel):
return tasks[2].result()
@util._check_socket
async def upload(self, source, target, name=None, timeout=None):
'''
Upload a stream to a device.
@@ -181,10 +203,12 @@ class Files(tunnel.Tunnel):
source (io.IOBase): An IO instance from which to read the data. Must be open for reading.
target (str): Path which to upload stream to on remote device
name (str): Pass if target points at a directory instead of the file path. In that case, this will be the name of the file.
timeout (int): duration in seconds to wait for a response before throwing an error
Raises:
:py:class:`~meshctrl.exceptions.FileTransferError`: File transfer failed. Info available on the `stats` property
:py:class:`~meshctrl.exceptions.FileTransferCancelled`: File transfer cancelled. Info available on the `stats` property
asyncio.TimeoutError: Command timed out
Returns:
dict: {result: bool whether upload succeeded, size: number of bytes uploaded}
@@ -198,17 +222,26 @@ class Files(tunnel.Tunnel):
raise request["error"]
return request["return"]
async def download(self, source, target, timeout=None):
def _http_download(self, url, target, timeout):
response = self._http_opener.open(url, timeout=timeout)
shutil.copyfileobj(response, target)
@util._check_socket
async def download(self, source, target, skip_http_attempt=False, skip_ws_attempt=False, timeout=None):
'''
Download a file from a device into a writable stream.
Args:
source (str): Path from which to download from device
target (io.IOBase): Stream to which to write data. If None, create new BytesIO which is both readable and writable.
skip_http_attempt (bool): Meshcentral has a way to download files through http(s) instead of through the websocket. This method tends to be much faster than using the websocket, so we try it first. Setting this to True will skip that attempt and just use the established websocket connection.
skip_ws_attempt (bool): Like skip_http_attempt, except just throw an error if the http attempt fails instead of trying with the websocket
timeout (int): duration in seconds to wait for a response before throwing an error
Raises:
:py:class:`~meshctrl.exceptions.FileTransferError`: File transfer failed. Info available on the `stats` property
:py:class:`~meshctrl.exceptions.FileTransferCancelled`: File transfer cancelled. Info available on the `stats` property
asyncio.TimeoutError: Command timed out
Returns:
dict: {result: bool whether download succeeded, size: number of bytes downloaded}
@@ -216,6 +249,29 @@ class Files(tunnel.Tunnel):
request_id = f"download_{self._get_request_id()}"
data = { "action": 'download', "sub": 'start', "id": request_id, "path": source }
request = {"id": request_id, "data": data, "type": "download", "source": source, "target": target, "size": 0, "finished": asyncio.Event(), "errored": asyncio.Event(), "error": None}
if not skip_http_attempt:
start_pos = target.tell()
try:
params = urllib.parse.urlencode({
"c": self._authcookie["cookie"],
"m": self._node.mesh.meshid.split("/")[-1],
"n": self._node.nodeid.split("/")[-1],
"f": source
})
url = self._session.url.replace('/control.ashx', f"/devicefile.ashx?{params}")
url = url.replace("wss://", "https://").replace("ws://", "http://")
loop = asyncio.get_event_loop()
await loop.run_in_executor(None, self._http_download, url, target, timeout)
size = target.tell() - start_pos
return {"result": True, "size": size}
except* Exception as eg:
if skip_ws_attempt:
size = target.tell() - start_pos
excs = eg.exceptions + (exceptions.FileTransferError("Errored", {"result": False, "size": size}),)
raise ExceptionGroup("File download failed", excs)
target.seek(start_pos)
await self._request_queue.put(request)
await asyncio.wait_for(request["finished"].wait(), timeout)
if request["error"] is not None:
@@ -230,7 +286,7 @@ class Files(tunnel.Tunnel):
return
if cmd["reqid"] == self._current_request["id"]:
if cmd["action"] == "uploaddone":
self._current_request["return"] = {"result": "success", "size": self._current_request["size"]}
self._current_request["return"] = {"result": True, "size": self._current_request["size"]}
self._current_request["finished"].set()
elif cmd["action"] == "uploadstart":
while True:
@@ -252,7 +308,7 @@ class Files(tunnel.Tunnel):
if self._current_request["inflight"] == 0 and self._current_request["complete"]:
await self._message_queue.put(json.dumps({ "action": 'uploaddone', "reqid": self._current_request["id"]}))
elif cmd["action"] == "uploaderror":
self._current_request["return"] = {"result": "canceled", "size": self._current_request["size"]}
self._current_request["return"] = {"result": False, "size": self._current_request["size"]}
self._current_request["error"] = exceptions.FileTransferError("Errored", self._current_request["return"])
self._current_request["errored"].set()
self._current_request["finished"].set()
@@ -268,7 +324,7 @@ class Files(tunnel.Tunnel):
self._current_request["target"].write(data[4:])
self._current_request["size"] += len(data)-4
if (data[3] & 1) != 0:
self._current_request["return"] = {"result": "success", "size": self._current_request["size"]}
self._current_request["return"] = {"result": True, "size": self._current_request["size"]}
self._current_request["finished"].set()
else:
await self._message_queue.put(json.dumps({ "action": 'download', "sub": 'ack', "id": self._current_request["id"] }))
@@ -279,7 +335,7 @@ class Files(tunnel.Tunnel):
if cmd["sub"] == "start":
await self._message_queue.put(json.dumps({ "action": 'download', "sub": 'startack', "id": self._current_request["id"] }))
elif cmd["sub"] == "cancel":
self._current_request["return"] = {"result": "canceled", "size": self._current_request["size"]}
self._current_request["return"] = {"result": False, "size": self._current_request["size"]}
self._current_request["error"] = exceptions.FileTransferCancelled("Cancelled", self._current_request["return"])
self._current_request["errored"].set()
self._current_request["finished"].set()

View File

@@ -215,7 +215,6 @@ class Session(object):
return self._command_id
async def close(self):
# Dunno yet
self._main_loop_task.cancel()
try:
await self._main_loop_task
@@ -1818,15 +1817,16 @@ class Session(object):
raise ValueError("No user or session given")
await self._message_queue.put(json.dumps({"action": "interuser", "data": data, "sessionid": session, "userid": user}))
async def upload(self, nodeid, source, target, unique_file_tunnel=False, timeout=None):
async def upload(self, node, source, target, unique_file_tunnel=False, timeout=None):
'''
Upload a stream to a device. This creates an _File and destroys it every call. If you need to upload multiple files, use {@link Session#file_explorer} instead.
Args:
nodeid (str): Unique id to upload stream to
node (~meshctrl.device.Device|str): Device or id of device to which to upload the file. If it is a device, it must have a ~meshctrl.mesh.Mesh device associated with it (the default). If it is a string, the device will be fetched prior to tunnel creation.
source (io.IOBase): An IO instance from which to read the data. Must be open for reading.
target (str): Path which to upload stream to on remote device
unique_file_tunnel (bool): True: Create a unique :py:class:`~meshctrl.files.Files` for this call, which will be cleaned up on return, else use cached or cache :py:class:`~meshctrl.files.Files`
timeout (int): duration in seconds to wait for a response before throwing an error
Raises:
:py:class:`~meshctrl.exceptions.FileTransferError`: File transfer failed. Info available on the `stats` property
@@ -1835,23 +1835,26 @@ class Session(object):
Returns:
dict: {result: bool whether upload succeeded, size: number of bytes uploaded}
'''
if not isinstance(node, device.Device):
node = await self.device_info(node)
if unique_file_tunnel:
async with self.file_explorer(nodeid) as files:
async with self.file_explorer(node) as files:
return await files.upload(source, target)
else:
files = await self._cached_file_explorer(nodeid, nodeid)
files = await self._cached_file_explorer(node, node.nodeid)
return await files.upload(source, target, timeout=timeout)
async def upload_file(self, nodeid, filepath, target, unique_file_tunnel=False, timeout=None):
async def upload_file(self, node, filepath, target, unique_file_tunnel=False, timeout=None):
'''
Friendly wrapper around :py:class:`~meshctrl.session.Session.upload` to upload from a filepath. Creates a ReadableStream and calls upload.
Args:
nodeid (str): Unique id to upload file to
node (~meshctrl.device.Device|str): Device or id of device to which to upload the file. If it is a device, it must have a ~meshctrl.mesh.Mesh device associated with it (the default). If it is a string, the device will be fetched prior to tunnel creation.
filepath (str): Path from which to read the data
target (str): Path which to upload file to on remote device
unique_file_tunnel (bool): True: Create a unique :py:class:`~meshctrl.files.Files` for this call, which will be cleaned up on return, else use cached or cache :py:class:`~meshctrl.files.Files`
timeout (int): duration in seconds to wait for a response before throwing an error
Raises:
:py:class:`~meshctrl.exceptions.FileTransferError`: File transfer failed. Info available on the `stats` property
@@ -1861,17 +1864,20 @@ class Session(object):
dict: {result: bool whether upload succeeded, size: number of bytes uploaded}
'''
with open(filepath, "rb") as f:
return await self.upload(nodeid, f, target, unique_file_tunnel, timeout=timeout)
return await self.upload(node, f, target, unique_file_tunnel, timeout=timeout)
async def download(self, nodeid, source, target=None, unique_file_tunnel=False, timeout=None):
async def download(self, node, source, target=None, skip_http_attempt=False, skip_ws_attempt=False, unique_file_tunnel=False, timeout=None):
'''
Download a file from a device into a writable stream. This creates an :py:class:`~meshctrl.files.Files` and destroys it every call. If you need to upload multiple files, use :py:class:`~meshctrl.session.Session.file_explorer` instead.
Args:
nodeid (str): Unique id to download file from
node (~meshctrl.device.Device|str): Device or id of device from which to download the file. If it is a device, it must have a ~meshctrl.mesh.Mesh device associated with it (the default). If it is a string, the device will be fetched prior to tunnel creation.
source (str): Path from which to download from device
target (io.IOBase): Stream to which to write data. If None, create new BytesIO which is both readable and writable.
skip_http_attempt (bool): Meshcentral has a way to download files through http(s) instead of through the websocket. This method tends to be much faster than using the websocket, so we try it first. Setting this to True will skip that attempt and just use the established websocket connection.
skip_ws_attempt (bool): Like skip_http_attempt, except just throw an error if the http attempt fails instead of trying with the websocket
unique_file_tunnel (bool): True: Create a unique :py:class:`~meshctrl.files.Files` for this call, which will be cleaned up on return, else use cached or cache :py:class:`~meshctrl.files.Files`
timeout (int): duration in seconds to wait for a response before throwing an error
Raises:
:py:class:`~meshctrl.exceptions.FileTransferError`: File transfer failed. Info available on the `stats` property
@@ -1880,29 +1886,34 @@ class Session(object):
Returns:
io.IOBase: The stream which has been downloaded into. Cursor will be at the beginning of where the file is downloaded.
'''
if not isinstance(node, device.Device):
node = await self.device_info(node)
if target is None:
target = io.BytesIO()
start = target.tell()
if unique_file_tunnel:
async with self.file_explorer(nodeid) as files:
async with self.file_explorer(node) as files:
await files.download(source, target)
target.seek(start)
return target
else:
files = await self._cached_file_explorer(nodeid, nodeid)
files = await self._cached_file_explorer(node, node.nodeid)
await files.download(source, target, timeout=timeout)
target.seek(start)
return target
async def download_file(self, nodeid, source, filepath, unique_file_tunnel=False, timeout=None):
async def download_file(self, node, source, filepath, skip_http_attempt=False, skip_ws_attempt=False, unique_file_tunnel=False, timeout=None):
'''
Friendly wrapper around :py:class:`~meshctrl.session.Session.download` to download to a filepath. Creates a WritableStream and calls download.
Args:
nodeid (str): Unique id to download file from
node (~meshctrl.device.Device|str): Device or id of device from which to download the file. If it is a device, it must have a ~meshctrl.mesh.Mesh device associated with it (the default). If it is a string, the device will be fetched prior to tunnel creation.
source (str): Path from which to download from device
filepath (str): Path to which to download data
skip_http_attempt (bool): Meshcentral has a way to download files through http(s) instead of through the websocket. This method tends to be much faster than using the websocket, so we try it first. Setting this to True will skip that attempt and just use the established websocket connection.
skip_ws_attempt (bool): Like skip_http_attempt, except just throw an error if the http attempt fails instead of trying with the websocket
unique_file_tunnel (bool): True: Create a unique :py:class:`~meshctrl.files.Files` for this call, which will be cleaned up on return, else use cached or cache :py:class:`~meshctrl.files.Files`
timeout (int): duration in seconds to wait for a response before throwing an error
Raises:
:py:class:`~meshctrl.exceptions.FileTransferError`: File transfer failed. Info available on the `stats` property
@@ -1912,24 +1923,39 @@ class Session(object):
None
'''
with open(filepath, "wb") as f:
await self.download(nodeid, source, f, unique_file_tunnel, timeout=timeout)
await self.download(node, source, f, unique_file_tunnel, timeout=timeout)
async def _cached_file_explorer(self, nodeid, _id):
async def _cached_file_explorer(self, node, _id):
if (_id not in self._file_tunnels or not self._file_tunnels[_id].alive):
self._file_tunnels[_id] = self.file_explorer(nodeid)
self._file_tunnels[_id] = await self.file_explorer(node).__aenter__()
await self._file_tunnels[_id].initialized.wait()
return self._file_tunnels[_id]
def file_explorer(self, nodeid):
def file_explorer(self, node):
'''
Create, initialize, and return an :py:class:`~meshctrl.files.Files` object for the given node
Args:
nodeid (str): Unique id on which to open file explorer
node (~meshctrl.device.Device|str): Device or id of device on which to open file explorer. If it is a device, it must have a ~meshctrl.mesh.Mesh device associated with it (the default). If it is a string, the device will be fetched prior to tunnel creation.
Returns:
:py:class:`~meshctrl.files.Files`: A newly initialized file explorer.
'''
return files.Files(self, nodeid)
'''
return _FileExplorerWrapper(self, node)
# This is a little yucky, but I can't get a good API otherwise. Since Tunnel objects are only useable as context managers anyway, this should be fine.
class _FileExplorerWrapper:
def __init__(self, session, node):
self.session = session
self.node = node
self._files = None
async def __aenter__(self):
if not isinstance(self.node, device.Device):
self.node = await self.session.device_info(self.node)
self._files = files.Files(self.session, self.node)
return await self._files.__aenter__()
async def __aexit__(self, exc_t, exc_v, exc_tb):
return await self._files.__aexit__(exc_t, exc_v, exc_tb)

View File

@@ -78,7 +78,6 @@ class Shell(tunnel.Tunnel):
read_bytes = 0
while True:
d = self._buffer.read1(length-read_bytes if length is not None else -1)
# print(f"read: {d}")
read_bytes += len(d)
ret.append(d)
if length is not None and read_bytes >= length:
@@ -163,7 +162,6 @@ class SmartShell(object):
command += "\n"
await self._shell.write(command)
data = await self._shell.expect(self._regex, timeout=timeout)
print(repr(data))
return data[:self._compiled_regex.search(data).span()[0]]
@property
@@ -178,14 +176,18 @@ class SmartShell(object):
def initialized(self):
return self._shell.initialized
@property
def _socket_open(self):
return self._shell._socket_open
async def close(self):
await self._init_task
await asyncio.wait_for(self._init_task, 10)
return await self._shell.close()
async def __aenter__(self):
await self._init_task
await self._shell.__aenter__()
await asyncio.wait_for(self._init_task, 10)
return self
async def __aexit__(self, *args):
return await self._shell.__aexit__(*args)
await self.close()

View File

@@ -27,6 +27,11 @@ class Tunnel(object):
self._message_queue = asyncio.Queue()
self._send_task = None
self._listen_task = None
self._ssl_context = None
if self._session._ignore_ssl:
self._ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
self._ssl_context.check_hostname = False
self._ssl_context.verify_mode = ssl.CERT_NONE
async def close(self):
self._main_loop_task.cancel()
@@ -45,22 +50,18 @@ class Tunnel(object):
async def _main_loop(self):
try:
authcookie = await self._session._send_command_no_response_id({ "action":"authcookie" })
self._authcookie = await self._session._send_command_no_response_id({ "action":"authcookie" })
options = {}
if self._session._ignore_ssl:
ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
ssl_context.check_hostname = False
ssl_context.verify_mode = ssl.CERT_NONE
options = { "ssl": ssl_context }
if self._ssl_context is not None:
options = { "ssl": self._ssl_context }
if (self.node_id.split('/') != 3) and (self._session._currentDomain is not None):
self.node_id = f"node/{self._session._currentDomain}/{self.node_id}"
if (len(self.node_id.split('/')) != 3):
self.node_id = f"node/{self._session._currentDomain or ""}/{self.node_id}"
self._tunnel_id = util._get_random_hex(6)
initialize_tunnel_response = await self._session._send_command({ "action": 'msg', "nodeid": self.node_id, "type": 'tunnel', "usage": 1, "value": '*/meshrelay.ashx?p=' + str(self._protocol) + '&nodeid=' + self.node_id + '&id=' + self._tunnel_id + '&rauth=' + authcookie["rcookie"] }, "initialize_tunnel")
initialize_tunnel_response = await self._session._send_command({ "action": 'msg', "nodeid": self.node_id, "type": 'tunnel', "usage": 1, "value": '*/meshrelay.ashx?p=' + str(self._protocol) + '&nodeid=' + self.node_id + '&id=' + self._tunnel_id + '&rauth=' + self._authcookie["rcookie"] }, "initialize_tunnel")
if initialize_tunnel_response.get("result", None) != "OK":
self._main_loop_error = exceptions.ServerError(initialize_tunnel_response.get("result", "Failed to initialize remote tunnel"))
self._socket_open.clear()
@@ -68,7 +69,8 @@ class Tunnel(object):
self.initialized.set()
return
self.url = self._session.url.replace('/control.ashx', '/meshrelay.ashx?browser=1&p=' + str(self._protocol) + '&nodeid=' + self.node_id + '&id=' + self._tunnel_id + '&auth=' + authcookie["cookie"])
self.url = self._session.url.replace('/control.ashx', '/meshrelay.ashx?browser=1&p=' + str(self._protocol) + '&nodeid=' + self.node_id + '&id=' + self._tunnel_id + '&auth=' + self._authcookie["cookie"])
async for websocket in util.proxy_connect(self.url, proxy_url=self._session._proxy, process_exception=util._process_websocket_exception, **options):
self.alive = True

View File

@@ -149,7 +149,7 @@ def _check_socket(f):
finally:
if not self.alive and self._main_loop_error is not None:
raise self._main_loop_error
elif not self.alive:
elif not self.alive and self.initialized.is_set():
raise exceptions.SocketError("Socket Closed")
return await f(self, *args, **kwargs)
return wrapper