client.py

from __future__ import annotations

import dataclasses
import logging
import sys
from pathlib import Path, PurePosixPath
from typing import TYPE_CHECKING

from mega import progress
from mega.api import APIContextManager
from mega.core import MegaCore
from mega.crypto import a32_to_base64, b64_to_a32, b64_url_encode, encrypt_attr, encrypt_key
from mega.data_structures import (
    AccountStats,
    Crypto,
    FileInfo,
    Node,
    NodeID,
    NodeType,
    UserResponse,
)
from mega.download import DownloadResults
from mega.errors import MegaNzError, RequestError, ValidationError
from mega.utils import Site, async_map, format_error, setup_logger

if TYPE_CHECKING:
    from contextlib import _GeneratorContextManager  # pyright: ignore[reportPrivateUsage]
    from os import PathLike

    import aiohttp
    import yarl

    from mega.filesystem import FileSystem, UserFileSystem

__all__ = ["MegaNzClient"]

logger = logging.getLogger(__name__)
_DOMAIN = Site.MEGA.value


class MegaNzClient(APIContextManager):
    """Interface with all the public methods of the API"""

    __slots__ = ("_core",)

    def __init__(self, session: aiohttp.ClientSession | None = None) -> None:
        super().__init__(session)
        self._core: MegaCore = MegaCore(self._api)
        if hasattr(sys, "ps1"):
            setup_logger(logging.DEBUG)

    def __repr__(self) -> str:
        return f"<{type(self).__name__}>(filesystem={self.cached_filesystem!r}, vault={self._core.vault!r}, logged_in={self.logged_in!r})"

    @property
    def cached_filesystem(self) -> UserFileSystem | None:
        return self._core.filesystem

    @property
    def logged_in(self) -> bool:
        return bool(self._core.vault.master_key)

    @property
    def progress_bar(self) -> _GeneratorContextManager[None]:
        return progress.new_progress()

    parse_url = staticmethod(MegaCore.parse_url)
    parse_file_url = staticmethod(MegaCore.parse_file_url)
    parse_folder_url = staticmethod(MegaCore.parse_folder_url)

    async def login(self, email: str | None = None, password: str | None = None) -> None:
        if self.logged_in:
            raise ValidationError("This client instance has already logged into an account")
        return await self._core.login(email, password)

    async def get_user(self) -> UserResponse:
        return await self._core.get_user()

    async def get_filesystem(self, *, force: bool = False) -> UserFileSystem:
        return await self._core.get_filesystem(force=force)

    async def search(self, query: str | PathLike[str], *, exclude_deleted: bool = True) -> dict[NodeID, PurePosixPath]:
        """Return nodes that have "query" as a substring on their path"""
        fs = await self.get_filesystem()
        return dict(fs.search(query, exclude_deleted=exclude_deleted))

    async def find(self, query: str | PathLike[str]) -> Node:
        """Return the single node located at *path*.

        NOTE: Mega's filesystem is **not POSIX-compliant**: multiple nodes may have the same path

        Raises `MultipleNodesFoundError` if more that one node has this path

        Raises `FileNotFoundError` if this path does not exists on the filesystem

        """
        fs = await self.get_filesystem()
        return fs.find(query)

    async def delete(self, node_id: NodeID) -> None:
        """Delete a file or folder by its private id (moves it to the trash bin)"""
        fs = await self.get_filesystem()
        await self.move(node_id, fs.trash_bin.id)

    async def destroy(self, node_id: NodeID) -> None:
        """Destroy a file or folder by its private id (bypasses trash bin)"""
        await self._core.destroy(node_id)

    async def empty_trash(self) -> None:
        """Deletes all file in the trash bin. Does nothing if the trash was already empty"""
        fs = await self.get_filesystem()
        trashed_files = [f.id for f in fs.deleted]
        if not trashed_files:
            return

        await self._core.destroy(*trashed_files)

    async def move(self, node_id: NodeID, target_id: NodeID) -> None:
        await self._core.move(node_id, target_id)

    async def add_contact(self, email: str) -> None:
        await self._core.edit_contact(email, add=True)

    async def remove_contact(self, email: str) -> None:
        await self._core.edit_contact(email, add=False)

    async def get_public_link(self, node: Node) -> str:
        if node.type is NodeType.FILE:
            public_key = a32_to_base64(node._crypto.full_key)
        elif node.type is NodeType.FOLDER:
            assert node._crypto.share_key
            public_key = a32_to_base64(node._crypto.share_key)

        else:
            msg = f"Can not get a public link for {node.type}"
            raise ValidationError(msg)

        public_handle = await self._core.public_handle_from_id(node.id)
        return f"{_DOMAIN}/{node.type.name.lower()}/{public_handle}#{public_key}"

    async def get_account_stats(self) -> AccountStats:
        resp = await self._core.get_account_stats()
        return AccountStats.parse(resp)

    async def export(self, node: Node) -> str:
        if node.type is NodeType.FILE:
            await self._core.export_file(node)
            return await self.get_public_link(node)

        if node.type is not NodeType.FOLDER:
            msg = f"Can only export files or folders, not {node.type}"
            raise ValidationError(msg)

        try:
            # If already exported
            return await self.get_public_link(node)
        except (RequestError, KeyError):
            await self._core.export_folder(node)
            fs = await self.get_filesystem(force=True)
            return await self.get_public_link(fs[node.id])

    async def get_public_filesystem(self, public_handle: NodeID, public_key: str) -> FileSystem:
        return await self._core.get_public_filesystem(public_handle, public_key)

    async def download(self, node: Node, output_dir: str | PathLike[str] | None = None) -> Path:
        """Download a file by it's file object."""
        file_info = await self._core.request_file_info(node.id)
        output_path = Path(output_dir or ".") / node.attributes.name
        return await self._core.download_file(file_info, node._crypto, output_path)

    async def download_url(self, url: str | yarl.URL, output_dir: str | PathLike[str] | None = None) -> DownloadResults:
        """Download a public file or folder by it's URL."""
        info = self.parse_url(url)
        if not info.is_folder:
            try:
                result = await self.download_public_file(info.public_handle, info.public_key, output_dir)
            except Exception as exc:
                result = exc
            return DownloadResults.split({info.public_handle: result})

        return await self.download_public_folder(info.public_handle, info.public_key, output_dir, info.selected_node)

    async def download_public_file(
        self,
        public_handle: NodeID,
        public_key: str,
        output_dir: str | PathLike[str] | None = None,
    ) -> Path:
        full_key = b64_to_a32(public_key)
        crypto = Crypto.decompose(full_key)
        file_info = await self._core.request_file_info(public_handle, is_public=True)
        output_name = self._core.decrypt_attrs(file_info._at, crypto.key).name
        output_path = Path(output_dir or ".") / output_name
        return await self._core.download_file(file_info, crypto, output_path)

    async def download_public_folder(
        self,
        public_handle: NodeID,
        public_key: str,
        output_dir: str | PathLike[str] | None = None,
        root_id: NodeID | None = None,
    ) -> DownloadResults:
        """Recursively download all files from a public folder, preserving its internal directory structure."""
        fs = await self.get_public_filesystem(public_handle, public_key)

        base_path = Path(output_dir or ".")
        folder_url = f"{_DOMAIN}/folder/{public_handle}#{public_key}"

        async def download(file: Node) -> tuple[NodeID, Path | Exception]:
            web_url = folder_url + f"/file/{file.id}"
            output_path = base_path / fs.relative_path(file.id)
            try:
                file_info = await self._core.request_file_info(file.id, public_handle)
                result = await self._core.download_file(file_info, file._crypto, output_path)

            except Exception as exc:
                msg = format_error(exc)
                logger.error(f'Unable to download {web_url} to "{output_path}" {msg}')
                result = exc

            return file.id, result

        results = await async_map(download, fs.files_from(root_id))
        return DownloadResults.split(dict(results))

    async def upload(self, file_path: str | PathLike[str], dest_node_id: NodeID | None = None) -> Node:
        if not dest_node_id:
            dest_node_id = (await self.get_filesystem()).root.id

        resp = await self._core.upload(file_path, dest_node_id)
        return self._core._deserialize_node(resp["f"][0])

    async def create_folder(self, path: str | PathLike[str]) -> Node:
        """Create a folder at the given path.

        If a folder already exists at that path, the existing folder node is
        returned instead of creating a new one.

        Note:
            Although MEGA allows creating nodes that share the same path,
            mega.py intentionally does **not** support this behaviour.
        """
        path = PurePosixPath(path).as_posix()
        fs = await self.get_filesystem()
        try:
            node = fs.find(path)

            if node.type is not NodeType.FOLDER:
                msg = f'Can not create a folder at "{path!s}". Node {node.id} already has this path'
                raise MegaNzError(msg)
            return node
        except FileNotFoundError:
            pass

        return await self._core.mkdir(path, fs.root.id)

    async def rename(self, node: Node, new_name: str) -> None:
        new_attrs = dataclasses.replace(node.attributes, name=new_name)
        attribs = b64_url_encode(encrypt_attr(new_attrs.serialize(), node._crypto.key))
        encrypted_key = a32_to_base64(encrypt_key(node._crypto.key, self._core.vault.master_key))

        _ = await self._core.api.post(
            {
                "a": "a",
                "attr": attribs,
                "key": encrypted_key,
                "n": node.id,
                "i": self._core.api.client_id,
            },
        )
        self._core.clear_cache()

    async def get_public_file_info(self, public_handle: NodeID, public_key: str) -> FileInfo:
        full_key = b64_to_a32(public_key)
        key = Crypto.decompose(full_key).key
        file_info = await self._core.request_file_info(public_handle, is_public=True)
        name = self._core.decrypt_attrs(file_info._at, key).name
        return dataclasses.replace(file_info, name=name)

    async def import_public_file(
        self,
        public_handle: NodeID,
        public_key: str,
        dest_node_id: NodeID | None = None,
    ) -> Node:
        """Import the public file into user account"""
        if not dest_node_id:
            dest_node_id = (await self.get_filesystem()).root.id

        return await self._core.import_file(public_handle, public_key, dest_node_id)