pylibmeshctrl/src/meshctrl/files.py

from . import tunnel
from . import constants
from . import exceptions
from . import util
import asyncio
import json
import importlib
import importlib.util
import shutil

# import urllib
# import urllib.request
import urllib.parse
old_parse = urllib.parse
# Default proxy handler uses OS defined no_proxy in order to be helpful. This is unhelpful for our usecase. Monkey patch out proxy getting functions, but don't effect the user's urllib instance.
spec = importlib.util.find_spec('urllib')
urllib = importlib.util.module_from_spec(spec)
spec.loader.exec_module(urllib)
spec = importlib.util.find_spec('urllib.request')
urllib.request = importlib.util.module_from_spec(spec)
spec.loader.exec_module(urllib.request)
urllib.parse = old_parse
urllib.request.getproxies_environment = lambda: {}
urllib.request.getproxies_registry = lambda: {}
urllib.request.getproxies_macosx_sysconf = lambda: {}
urllib.request.getproxies = lambda: {}

class Files(tunnel.Tunnel):
    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()
        self._download_finished.set()
        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": self._session._proxy,
                "https": self._session._proxy,
                "no": ""
            }
        self._proxy_handler = urllib.request.ProxyHandler(proxies=proxies)
        self._http_opener = urllib.request.build_opener(self._proxy_handler, urllib.request.HTTPSHandler(context=self._session._ssl_context))


    def _get_request_id(self):
        self._request_id = (self._request_id+1)%(2**32-1)
        return self._request_id

    async def close(self):
        self._handle_requests_task.cancel()
        try:
            await self._handle_requests_task
        except asyncio.CancelledError:
            pass
        await super().close()

    async def _handle_requests(self):
        try:
            while True:
                request = await self._request_queue.get()
                self._current_request = request
                self._download_finished = request["finished"]
                await self._message_queue.put(json.dumps(request["data"]))
                await request["finished"].wait()
                self._current_request = None

        except asyncio.CancelledError:
            while True:
                try:
                    request = self._request_queue.get_nowait()
                    request["error"] = exceptions.SocketError("Socket Closed")
                    request["errored"].set()
                    request["finished"].set()
                except asyncio.QueueEmpty:
                    break
            raise
            

    @util._check_socket
    async def _send_command(self, data, name, timeout=None):
        request_id = f"meshctrl_{name}_{self._get_request_id()}"
        request = {"id": request_id, "data": data, "return": None, "type": name, "finished": asyncio.Event(), "errored":asyncio.Event(), "error": None}
        await self._request_queue.put(request)

        await asyncio.wait_for(request["finished"].wait(), timeout=timeout)
        if request["error"] is not None:
            raise request["error"]
        return request["return"]

    async def ls(self, directory, timeout=None):
        """
        Return a directory listing from the device

        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

        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"]

    async def _listen_for_pass(self, tasks):
        async for event in self._session.events({"event": {"etype": "node", "action": "agentlog"}}):
            if not event["event"]["msg"].startswith("Started"):
                self._current_request["return"] = event["event"]["msg"]
                self._current_request["finished"].set()
                tasks[1].cancel()
                break

    async def _listen_for_error(self, tasks):
        async for event in self._session.events({"action":"msg", "type":"console"}):
            self._current_request["error"] = exceptions.ServerError(event["value"])
            self._current_request["errored"].set()
            self._current_request["finished"].set()
            tasks[0].cancel()
            break

    async def mkdir(self, directory, timeout=None):
        """
        Create a directory on the device

        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
        """
        tasks = []
        async with asyncio.TaskGroup() as tg:
            tasks.append(tg.create_task(asyncio.wait_for(self._listen_for_pass(tasks), timeout)))
            tasks.append(tg.create_task(asyncio.wait_for(self._listen_for_error(tasks), timeout)))
            tasks.append(tg.create_task(self._send_command({"action": "mkdir", "path": directory}, "mkdir", timeout=timeout)))



        return tasks[2].result().startswith("Create folder")

    async def rm(self, path, files, recursive=False, timeout=None):
        """
        Remove a set of files or directories from the device. This API doesn't error if the file doesn't exist.

        Args:
            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'.
        """
        if isinstance(files, str):
            files = [files]
        tasks = []

        async with asyncio.TaskGroup() as tg:
            tasks.append(tg.create_task(asyncio.wait_for(self._listen_for_pass(tasks), timeout)))
            tasks.append(tg.create_task(asyncio.wait_for(self._listen_for_error(tasks), timeout)))
            tasks.append(tg.create_task(self._send_command({"action": "rm", "delfiles": files, "rec": recursive, "path": path}, "rm", timeout=timeout)))


        return tasks[2].result()

    async def rename(self, path, name, new_name, timeout=None):
        """
        Rename a file or folder on the device. This API doesn't error if the file doesn't exist.

        Args:
            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"'.
        """
        tasks = []

        async with asyncio.TaskGroup() as tg:
            tasks.append(tg.create_task(asyncio.wait_for(self._listen_for_pass(tasks), timeout)))
            tasks.append(tg.create_task(asyncio.wait_for(self._listen_for_error(tasks), timeout)))
            tasks.append(tg.create_task(self._send_command({"action": "rename", "path": path, "oldname": name, "newname": new_name}, "rename", timeout=timeout)))


        return tasks[2].result()

    @util._check_socket
    async def upload(self, source, target, name=None, timeout=None):
        '''
        Upload a stream to a device.

        Args:
            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}
        '''
        request_id = f"upload_{self._get_request_id()}"
        data = { "action": 'upload', "reqid": request_id, "path": target, "name": name}
        request = {"id": request_id, "data": data, "type": "upload", "source": source, "target": target, "name": name, "size": 0, "complete": False, "inflight": 0, "finished": asyncio.Event(), "errored":asyncio.Event(), "error": None}
        await self._request_queue.put(request)
        await asyncio.wait_for(request["finished"].wait(), timeout)
        if request["error"] is not None:
            raise request["error"]
        return request["return"]

    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}
        '''
        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:
            raise request["error"]
        return request["return"]

    async def _handle_upload(self, data):
        cmd = None
        try:
            cmd = json.loads(data)
        except:
            return
        if cmd["reqid"] == self._current_request["id"]:
            if cmd["action"] == "uploaddone":
                self._current_request["return"] = {"result": True, "size": self._current_request["size"]}
                self._current_request["finished"].set()
            elif cmd["action"] == "uploadstart":
                while True:
                    data = self._current_request["source"].read(self._chunk_size)
                    if len(data) == 0:
                        self._current_request["complete"] = True
                        if self._current_request["inflight"] == 0:
                            await self._message_queue.put(json.dumps({ "action": 'uploaddone', "reqid": self._current_request["id"]}))
                        break
                    else:
                        self._current_request["size"] += len(data)
                        if data[0] == 0 or data[0] == 123:
                            data = b'\0' + data
                        await self._message_queue.put(data)
                        self._current_request["inflight"] += 1
                    await asyncio.sleep(0)
            elif cmd["action"] == "uploadack":
                self._current_request["inflight"] -= 1
                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": 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()

    async def _handle_download(self, data):
        cmd = None
        try:
            cmd = json.loads(data)
        except:
            pass
        if cmd is None:
            if len(data) > 4:
                self._current_request["target"].write(data[4:])
                self._current_request["size"] += len(data)-4
            if (data[3] & 1) != 0:
                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"] }))
        else:
            if cmd["action"] == "download":
                if cmd["id"] != self._current_request["id"]:
                    return
                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": 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()

    async def _handle_action(self, data):
        self._current_request["return"] = json.loads(data)
        self._download_finished.set()


    async def _listen_data_task(self, websocket):
        async for message in websocket:
            if self.initialized.is_set():
                if message[0] == 123 and self._current_request is not None and self._current_request["type"] not in ("upload", "download"):
                    await self._handle_action(message)
                elif self._current_request is not None and self._current_request["type"] == "upload":
                    await self._handle_upload(message)
                elif self._current_request is not None and self._current_request["type"] == "download":
                    await self._handle_download(message)
            else:
                self.recorded = False
                if message == "cr":
                    self.recorded = True

                await self._message_queue.put(f"{self._protocol}".encode())
                self.alive = True
                self.initialized.set()