From a3c721318d6c6332cae66fd52f217addee7eb525 Mon Sep 17 00:00:00 2001
From: Josiah Baldwin <jbaldwin8889@gmail.com>
Date: Thu, 12 Dec 2024 16:06:18 -0800
Subject: [PATCH] Added ability to download files over http(s)

Also fixed some tests and a couple other bugs
---
 src/meshctrl/files.py                         | 70 +++++++++++++++++--
 src/meshctrl/session.py                       | 70 +++++++++++++------
 src/meshctrl/shell.py                         | 12 ++--
 src/meshctrl/tunnel.py                        | 24 ++++---
 src/meshctrl/util.py                          |  2 +-
 tests/.gitignore                              |  4 +-
 tests/environment/__init__.py                 |  6 +-
 .../{ => config}/meshcentral/data/config.json |  0
 tests/environment/meshcentral.dockerfile      |  2 +-
 tests/test_files.py                           | 20 +++++-
 tests/test_sanity.py                          |  3 +-
 tests/test_session.py                         | 21 +++---
 12 files changed, 166 insertions(+), 68 deletions(-)
 rename tests/environment/{ => config}/meshcentral/data/config.json (100%)

diff --git a/src/meshctrl/files.py b/src/meshctrl/files.py
index ffc659d..19a704e 100644
--- a/src/meshctrl/files.py
+++ b/src/meshctrl/files.py
@@ -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()
diff --git a/src/meshctrl/session.py b/src/meshctrl/session.py
index 22d8e9a..3beb0f1 100644
--- a/src/meshctrl/session.py
+++ b/src/meshctrl/session.py
@@ -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)
 
-        
\ No newline at end of file
+
+# 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)
\ No newline at end of file
diff --git a/src/meshctrl/shell.py b/src/meshctrl/shell.py
index d504495..46a4589 100644
--- a/src/meshctrl/shell.py
+++ b/src/meshctrl/shell.py
@@ -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()
diff --git a/src/meshctrl/tunnel.py b/src/meshctrl/tunnel.py
index 9a68c75..2f5278b 100644
--- a/src/meshctrl/tunnel.py
+++ b/src/meshctrl/tunnel.py
@@ -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
diff --git a/src/meshctrl/util.py b/src/meshctrl/util.py
index 532d3a8..fe36fc5 100644
--- a/src/meshctrl/util.py
+++ b/src/meshctrl/util.py
@@ -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
diff --git a/tests/.gitignore b/tests/.gitignore
index 0fd169d..8dbb45f 100644
--- a/tests/.gitignore
+++ b/tests/.gitignore
@@ -1,2 +1,2 @@
-data
-environment/scripts/meshcentral/users.json
\ No newline at end of file
+/data
+/environment/scripts/meshcentral/users.json
\ No newline at end of file
diff --git a/tests/environment/__init__.py b/tests/environment/__init__.py
index 1badb22..84aef33 100644
--- a/tests/environment/__init__.py
+++ b/tests/environment/__init__.py
@@ -43,10 +43,8 @@ class Agent(object):
         return self
 
     def __exit__(self, exc_t, exc_v, exc_tb):
-        try:
-            requests.post("{self._clienturl}/remove-agent/{self.nodeid}")
-        except:
-            pass
+        requests.post(f"{self._clienturl}/remove-agent/{self.nodeid}")
+
 
 class TestEnvironment(object):
     def __init__(self):
diff --git a/tests/environment/meshcentral/data/config.json b/tests/environment/config/meshcentral/data/config.json
similarity index 100%
rename from tests/environment/meshcentral/data/config.json
rename to tests/environment/config/meshcentral/data/config.json
diff --git a/tests/environment/meshcentral.dockerfile b/tests/environment/meshcentral.dockerfile
index 414dcc3..b2e213a 100644
--- a/tests/environment/meshcentral.dockerfile
+++ b/tests/environment/meshcentral.dockerfile
@@ -3,5 +3,5 @@ RUN apk add curl
 RUN apk add python3
 WORKDIR /opt/meshcentral/
 COPY ./scripts/meshcentral ./scripts
-COPY ./meshcentral/data /opt/meshcentral/meshcentral-data
+COPY ./config/meshcentral/data /opt/meshcentral/meshcentral-data
 CMD ["python3", "/opt/meshcentral/scripts/create_users.py"]
\ No newline at end of file
diff --git a/tests/test_files.py b/tests/test_files.py
index fb22c94..e69152a 100644
--- a/tests/test_files.py
+++ b/tests/test_files.py
@@ -5,6 +5,7 @@ import meshctrl
 import requests
 import io
 import random
+import time
 
 async def test_commands(env):
     async with meshctrl.Session("wss://" + env.dockerurl, user="admin", password=env.users["admin"], ignore_ssl=True, proxy=env.proxyurl) as admin_session:
@@ -78,7 +79,7 @@ async def test_upload_download(env):
                 async with admin_session.file_explorer(agent.nodeid) as files:
                     r = await files.upload(upfilestream, f"{pwd}/test", timeout=5)
                     print("\ninfo files_upload: {}\n".format(r))
-                    assert r["result"] == "success", "Upload failed"
+                    assert r["result"] == True, "Upload failed"
                     assert r["size"] == len(randdata), "Uploaded wrong number of bytes"
                     for f in await files.ls(pwd, timeout=5):
                         if f["n"] == "test" and f["t"] == meshctrl.constants.FileType.FILE:
@@ -95,10 +96,23 @@ async def test_upload_download(env):
                     else:
                         raise Exception("Uploaded file not found")
 
-                    r = await files.download(f"{pwd}/test", downfilestream, timeout=5)
+                    start = time.perf_counter()
+                    r = await files.download(f"{pwd}/test", downfilestream, skip_ws_attempt=True, timeout=5)
                     print("\ninfo files_download: {}\n".format(r))
-                    assert r["result"] == "success", "Domnload failed"
+                    assert r["result"] == True, "Domnload failed"
                     assert r["size"] == len(randdata), "Downloaded wrong number of bytes"
+                    print(f"http download time: {time.perf_counter()-start}")
+
+                    downfilestream.seek(0)
+                    assert downfilestream.read() == randdata, "Got wrong data back"
+                    downfilestream.seek(0)
+
+                    start = time.perf_counter()
+                    r = await files.download(f"{pwd}/test", downfilestream, skip_http_attempt=True, timeout=5)
+                    print("\ninfo files_download: {}\n".format(r))
+                    assert r["result"] == True, "Domnload failed"
+                    assert r["size"] == len(randdata), "Downloaded wrong number of bytes"
+                    print(f"ws download time: {time.perf_counter()-start}")
 
                     downfilestream.seek(0)
                     assert downfilestream.read() == randdata, "Got wrong data back"
diff --git a/tests/test_sanity.py b/tests/test_sanity.py
index 9dbc4d9..835c876 100644
--- a/tests/test_sanity.py
+++ b/tests/test_sanity.py
@@ -43,7 +43,7 @@ async def test_urlparse():
     try:
         async with meshctrl.Session("wss://localhost", user="unprivileged", password="Not a real password", ignore_ssl=True) as s:
             pass
-    except* TimeoutError:
+    except* asyncio.TimeoutError:
         #We're not running a server, so timeout is our expected outcome
         pass
 
@@ -52,5 +52,4 @@ async def test_urlparse():
         async with meshctrl.Session("https://localhost", user="unprivileged", password="Not a real password", ignore_ssl=True) as s:
             pass
     except* ValueError:
-        #We're not running a server, so timeout is our expected outcome
         pass
\ No newline at end of file
diff --git a/tests/test_session.py b/tests/test_session.py
index 484ea6a..f88db24 100644
--- a/tests/test_session.py
+++ b/tests/test_session.py
@@ -29,7 +29,7 @@ async def test_admin(env):
             assert len(no_sessions.keys()) == 0, "non-admin has admin acess"
 
         assert len(admin_users) == len(env.users.keys()), "Admin cannot see correct number of users"
-        assert len(admin_sessions) == 2, "Admin cannot see correct number of oser sessions"
+        assert len(admin_sessions) == 2, "Admin cannot see correct number of user sessions"
 
 async def test_auto_reconnect(env):
     async with meshctrl.Session(env.mcurl, user="admin", password=env.users["admin"], ignore_ssl=True, auto_reconnect=True) as admin_session:
@@ -39,6 +39,7 @@ async def test_auto_reconnect(env):
 
     # As above, but with proxy
     async with meshctrl.Session("wss://" + env.dockerurl, user="admin", password=env.users["admin"], ignore_ssl=True, auto_reconnect=True, proxy=env.proxyurl) as admin_session:
+
         env.restart_mesh()
         for i in range(3):
             try:
@@ -284,14 +285,14 @@ async def test_mesh_device(env):
             except:
                 raise Exception("Failed to run command on device after it was moved to a new mesh while having individual device permissions")
 
-        r = await admin_session.remove_users_from_device_group((await privileged_session.user_info())["_id"], mesh.meshid, timeout=10)
-        print("\ninfo remove_users_from_device_group: {}\n".format(r))
-        assert (await admin_session.remove_users_from_device(agent.nodeid, (await unprivileged_session.user_info())["_id"], timeout=10)), "Failed to remove user from device"
-
-        assert (r[(await privileged_session.user_info())["_id"]]["success"]), "Failed to remove user from devcie group"
+            r = await admin_session.remove_users_from_device_group((await privileged_session.user_info())["_id"], mesh.meshid, timeout=10)
+            print("\ninfo remove_users_from_device_group: {}\n".format(r))
+            assert (r[(await privileged_session.user_info())["_id"]]["success"]), "Failed to remove user from device group"
+            assert (await admin_session.remove_users_from_device(agent.nodeid, (await unprivileged_session.user_info())["_id"], timeout=10)), "Failed to remove user from device"
+            
 
         assert (await admin_session.remove_device_group(mesh.meshid, timeout=10)), "Failed to remove device group"
-        assert (await admin_session.remove_device_group(mesh2.name, isname=True, timeout=10)), "Failed to remove device group"
+        assert (await admin_session.remove_device_group(mesh2.name, isname=True, timeout=10)), "Failed to remove device group by name"
         assert not (await admin_session.add_users_to_device_group((await privileged_session.user_info())["_id"], mesh.meshid, rights=meshctrl.constants.MeshRights.fullrights, timeout=5))[(await privileged_session.user_info())["_id"]]["success"], "Added user to device group which doesn't exist?"
 
 async def test_user_groups(env):
@@ -416,12 +417,12 @@ async def test_session_files(env):
 
                 r = await admin_session.upload(agent.nodeid, upfilestream, f"{pwd}/test", timeout=5)
                 print("\ninfo files_upload: {}\n".format(r))
-                assert r["result"] == "success", "Upload failed"
+                assert r["result"] == True, "Upload failed"
                 assert r["size"] == len(randdata), "Uploaded wrong number of bytes"
 
                 r = await admin_session.upload_file(agent.nodeid, os.path.join(thisdir, "data", "test"), f"{pwd}/test2", timeout=5)
                 print("\ninfo files_upload: {}\n".format(r))
-                assert r["result"] == "success", "Upload failed"
+                assert r["result"] == True, "Upload failed"
                 assert r["size"] == len(randdata), "Uploaded wrong number of bytes"
 
                 s = await admin_session.download(agent.nodeid, f"{pwd}/test", timeout=5)
@@ -437,7 +438,7 @@ async def test_session_files(env):
 
                 r = await admin_session.upload_file(agent.nodeid, os.path.join(thisdir, "data", "test"), f"{pwd}/test2", unique_file_tunnel=True, timeout=5)
                 
-                assert r["result"] == "success", "Upload failed"
+                assert r["result"] == True, "Upload failed"
                 assert r["size"] == len(randdata), "Uploaded wrong number of bytes"
 
                 await admin_session.download_file(agent.nodeid, f"{pwd}/test2", os.path.join(thisdir, "data", "test"), unique_file_tunnel=True, timeout=5)