blender-package-manager-addon/bpkg_manager/subproc.py

"""
All the stuff that needs to run in a subprocess.
"""

import logging
import pathlib
import shutil
import json


class Message:
    """Superclass for all message sent over pipes."""


class BlenderMessage(Message):
    """Superclass for all messages sent from Blender to the subprocess."""


class SubprocMessage(Message):
    """Superclass for all messages sent from the subprocess to Blender."""


class Abort(BlenderMessage):
    """Sent when the user requests abortion of a task."""


class Progress(SubprocMessage):
    """Send from subprocess to Blender to report progress.

    :ivar progress: the progress percentage, from 0-1.
    """

    def __init__(self, progress: float):
        self.progress = progress

class SubprocError(SubprocMessage):
    """Superclass for all fatal error messages sent from the subprocess."""

    def __init__(self, message: str):
        self.message = message

class SubprocWarning(SubprocMessage):
    """Superclass for all non-fatal warning messages sent from the subprocess."""

    def __init__(self, message: str):
        self.message = message

class InstallError(SubprocError):
    """Sent when there was an error installing something."""

class FileConflictError(InstallError):
    """Sent when installation would overwrite existing files."""

    def __init__(self, message: str, conflicts: list):
        self.message = message
        self.conflicts = conflicts

class DownloadError(SubprocMessage):
    """Sent when there was an error downloading something."""

    def __init__(self, status_code: int, description: str):
        self.status_code = status_code
        self.description = description


class Success(SubprocMessage):
    """Sent when an operation finished sucessfully."""

class Result(SubprocMessage):
    """Sent when an operation returns data to be used on the parent process."""

    def __init__(self, data):
        self.data = data

class Aborted(SubprocMessage):
    """Sent as response to Abort message."""


class InstallException(Exception):
    """Raised when there is an error during installation"""

class DownloadException(Exception):
    """Raised when there is an error downloading something"""

    def __init__(self, status_code: int, message: str):
        self.status_code = status_code
        self.message = message

class BadRepository(Exception):
    """Raised when reading a repository results in an error"""


class InplaceBackup:
    """Utility for moving a file out of the way by appending a '~'"""

    log = logging.getLogger('%s.inplace-backup' % __name__)

    def __init__(self, path: pathlib.Path):
        self.path = path
        self.backup()

    def backup(self):
        """Move 'path' to 'path~'"""
        self.backup_path = pathlib.Path(str(self.path) + '~')
        if self.backup_path.exists():
            self.log.warning("Overwriting existing backup '{}'".format(self.backup_path))
            self._rm(self.backup_path)

        shutil.move(str(self.path), str(self.backup_path))

    def restore(self):
        """Move 'path~' to 'path'"""
        if not self.backup_path:
            raise RuntimeError("Can't restore file before backing it up")

        if self.path.exists():
            self.log.warning("Overwriting '{0}' with backup file".format(self.path))
            self._rm(self.backup_path)

        shutil.move(str(self.backup_path), str(self.path))

    def remove(self):
        """Remove 'path~'"""
        self._rm(self.backup_path)

    def _rm(self, path: pathlib.Path):
        """Just delete whatever is specified by `path`"""
        if path.is_file():
            path.unlink()
        else:
            shutil.rmtree(str(path))



class Package:
    """
    Stores package methods and metadata
    """

    log = logging.getLogger(__name__ + ".Repository")

    def __init__(self, package_dict:dict = None):
        self.from_dict(package_dict)

    def to_dict(self) -> dict:
        """
        Return a dict representation of the package
        """
        return {
                'bl_info': self.bl_info,
                'url': self.url,
                }

    def from_dict(self, package_dict: dict):
        """
        Get attributes from a dict such as produced by `to_dict`
        """
        if package_dict is None:
            package_dict = {}

        for attr in ('name', 'url', 'bl_info'):
            setattr(self, attr, package_dict.get(attr))



class Repository:
    """
    Stores repository metadata (including packages)
    """

    log = logging.getLogger(__name__ + ".Repository")

    def __init__(self, url=None):
        self.set_from_dict({'url': url})
        self.log.debug("Initializing repository: %s", self.to_dict())

    def refresh(self):
        """
        Requests repo.json from URL and embeds etag/last-modification headers
        """
        import requests

        if self.url is None:
            raise ValueError("Cannot refresh repository without a URL")

        self.log.debug("Refreshing repository from %s", self.url)

        req_headers = {}
        # Do things this way to avoid adding empty objects/None to the req_headers dict
        if self._headers:
            try:
                req_headers['If-None-Match'] = self._headers['etag']
            except KeyError:
                pass
            try:
                req_headers['If-Modified-Since'] = self._headers['last-modified']
            except KeyError:
                pass

        resp = requests.get(self.url, headers=req_headers)

        try:
            resp.raise_for_status()
        except requests.HTTPError as err:
            self.log.error('Error downloading %s: %s', self.url, err)
            raise DownloadException(resp.status_code, resp.reason) from err

        if resp.status_code == requests.codes.not_modified:
            self.log.debug("Packagelist not modified")
            return

        resp_headers = {}
        try:
            resp_headers['etag'] = resp.headers['etag']
        except KeyError:
            pass
        try:
            resp_headers['last-modified'] = resp.headers['last-modified']
        except KeyError:
            pass

        self.log.debug("Found headers: %s", resp_headers)

        
        repodict = resp.json()
        repodict['_headers'] = resp_headers

        self.set_from_dict(repodict)


    def to_dict(self) -> dict:
        """
        Return a dict representation of the repository
        """
        self.log.debug("Rendering to a dict")
        self.log.debug("url: %s", self.url)
        return {
                'name':     self.name,
                'packages': [p.to_dict() for p in self.packages] if self.packages is not None else None,
                'url':      self.url,
                '_headers': self._headers,
                }

    def set_from_dict(self, repodict: dict):
        """
        Get repository attributes from a dict such as produced by `to_dict`
        """
        if repodict is None:
            repodict = {}

        for attr in ('name', 'url', 'packages', '_headers'):
            if attr == 'packages':
                value = set(Package(pkg) for pkg in repodict.get('packages', []))
            else:
                value = repodict.get(attr)

            if value is None:
                try:
                    value = getattr(self, attr)
                except AttributeError:
                    pass

            setattr(self, attr, value)

    @classmethod
    def from_dict(cls, repodict: dict):
        """
        Like `set_from_dict`, but immutable
        """
        repo = cls()
        repo.set_from_dict(repodict)
        return repo

    def to_file(self, path: pathlib.Path):
        """
        Dump repository to a json file at `path`.
        """
        if self.packages is None:
            self.log.warning("Writing an empty repository")

        with path.open('w', encoding='utf-8') as repo_file:
            json.dump(self.to_dict(), repo_file, indent=4, sort_keys=True)
        self.log.debug("Repository written to %s" % path)

    @classmethod
    def from_file(cls, path: pathlib.Path):
        """
        Read repository from a json file at `path`.
        """
        try:
            repo_file = path.open('r', encoding='utf-8')
        except IOError as err:
            raise BadRepository from err

        with repo_file:
            try:
                repo = cls.from_dict(json.load(repo_file))
            except Exception as err:
                raise BadRepository from err

        cls.log.debug("Repository read from %s", path)
        return repo


def _download(pipe_to_blender, package_url: str, download_dir: pathlib.Path) -> pathlib.Path:
    """Downloads the given package

    :returns: path to the downloaded file, or None in case of error.
    """

    import requests

    log = logging.getLogger('%s.download' % __name__)
    log.info('Going to download %s to %s', package_url, download_dir)
    pipe_to_blender.send(Progress(0.0))

    log.info('Downloading %s', package_url)
    resp = requests.get(package_url, stream=True, verify=True)

    try:
        resp.raise_for_status()
    except requests.HTTPError as ex:
        log.error('Error downloading %s: %s', package_url, ex)
        pipe_to_blender.send(DownloadError(resp.status_code, str(ex)))
        return None

    try:
        # Use float so that we can also use infinity
        content_length = float(resp.headers['content-length'])
    except KeyError:
        log.warning('Server did not send content length, cannot report progress.')
        content_length = float('inf')

    # TODO: check if there's enough disk space.

    # TODO: get filename from Content-Disposition header, if available.
    # TODO: use urllib.parse to parse the URL.
    local_filename = package_url.split('/')[-1] or 'download.tmp'
    local_fpath = download_dir / local_filename

    downloaded_length = 0
    with local_fpath.open('wb') as outfile:
        for chunk in resp.iter_content(chunk_size=1024 ** 2):
            # Handle abort messages from Blender
            while pipe_to_blender.poll():
                recvd = pipe_to_blender.recv()
                if isinstance(recvd, Abort):
                    log.warning('Aborting download of %s by request', package_url)
                    pipe_to_blender.send(Aborted())
                    return None
                log.warning('Unknown message %s received, ignoring', recvd)

            if not chunk:  # filter out keep-alive new chunks
                continue

            outfile.write(chunk)

            downloaded_length += len(chunk)

            # TODO: use multiplier for progress, so that we can count up to 70% and
            # leave 30% "progress" for installation of the package.
            pipe_to_blender.send(Progress(downloaded_length / content_length))

    return local_fpath

def _install(pipe_to_blender, pkgpath: pathlib.Path, dest: pathlib.Path, searchpaths: list):
    """Extracts/moves package at `pkgpath` to `dest`"""
    import zipfile

    log = logging.getLogger('%s.install' % __name__)

    log.debug("Starting installation")
    pipe_to_blender.send(Progress(0.0))

    if not pkgpath.is_file():
        log.error("File to install isn't a file: %s", pkgpath)
        pipe_to_blender.send(InstallError("Package is not a file"))
        raise InstallException

    if not dest.is_dir():
        log.error("Destination for install isn't a directory: %s", dest)
        pipe_to_blender.send(InstallError("Destination is not a directory"))
        raise InstallException

    # TODO: check to make sure addon/package isn't already installed elsewhere

    # The following is adapted from `addon_install` in bl_operators/wm.py

    # check to see if the file is in compressed format (.zip)
    if zipfile.is_zipfile(pkgpath):
        log.debug("Package is zipfile")
        try:
            file_to_extract = zipfile.ZipFile(str(pkgpath), 'r')
        except Exception as err:
            # TODO HACK: see if it makes sense to make a single InstallError class which inherits from both Exception and SubprocError
            # It sounds weird, but it could avoid the need for duplication here; create a new InstallError with the right message,
            # then send it to blender and also raise it.
            pipe_to_blender.send(InstallError("Failed to read zip file: %s" % err))
            raise InstallException from err

        def root_files(filelist: list) -> list:
            """Some string parsing to get a list of the root contents of a zip from its namelist"""
            rootlist = []
            for f in filelist:
                # Get all names which have no path separators (root level files)
                # or have a single path separator at the end (root level directories).
                if len(f.rstrip('/').split('/')) == 1:
                    rootlist.append(f)
            return rootlist

        conflicts = [dest / f for f in root_files(file_to_extract.namelist()) if (dest / f).exists()]
        backups = []
        for conflict in conflicts:
            backups.append(InplaceBackup(conflict))

        try:
            file_to_extract.extractall(str(dest))
        except Exception as err:
            for backup in backups:
                backup.restore()
            pipe_to_blender.send(InstallError("Failed to extract zip file to '%s': %s" % (dest, err)))
            raise InstallException from err

        for backup in backups:
            backup.remove()

    else:
        log.debug("Package is pyfile")
        dest_file = (dest / pkgpath.name)

        if dest_file.exists():
            backup = InplaceBackup(dest_file)

        try:
            shutil.copyfile(str(pkgpath), str(dest_file))
        except Exception as err:
            backup.restore()
            pipe_to_blender.send(InstallError("Failed to copy file to '%s': %s" % (dest, err)))
            raise InstallException from err

    try:
        pkgpath.unlink()
        log.debug("Removed cached package: %s", pkgpath)
    except Exception as err:
        pipe_to_blender.send(SubprocWarning("Failed to remove package from cache: %s" % err))
        raise InstallException from err

    pipe_to_blender.send(Progress(1.0))
    return


def download_and_install(pipe_to_blender, package_url: str, install_path: pathlib.Path, search_paths: list):
    """Downloads and installs the given package."""

    from . import cache

    log = logging.getLogger('%s.download_and_install' % __name__)

    # log.debug(utils.user_resource('SCRIPTS', 'blorp'))

    cache_dir = cache.cache_directory('downloads')
    downloaded = _download(pipe_to_blender, package_url, cache_dir)

    if not downloaded:
        log.debug('Download failed/aborted, not going to install anything.')
        return

    # Only send success if _install doesn't throw an exception
    # Maybe not the best way to do this? (will also catch exceptions which occur during message sending)
    try:
        _install(pipe_to_blender, downloaded, install_path, search_paths)
        pipe_to_blender.send(Success())
    except InstallException as err:
        #TODO
        log.error("Failed to install package: %s", err)
        pipe_to_blender.send(InstallError(err))

def refresh(pipe_to_blender, storage_path: pathlib.Path, repository_url: str):
    """Retrieves and stores the given repository"""

    log = logging.getLogger(__name__ + '.refresh')

    repo_path = storage_path / 'repo.json'
    try:
        repo = Repository.from_file(repo_path)
    except BadRepository as err:
        log.warning("Failed to read existing repository: %s. Continuing download.", err)
        repo = Repository(repository_url)

    if repo.url != repository_url:
        # We're getting a new repository
        repo = Repository(repository_url)

    try:
        repo.refresh()
    except DownloadException as err:
        pipe_to_blender.send(DownloadError(err.status_code, err.message))

    repo.to_file(repo_path) # TODO: this always writes even if repo wasn't changed
    pipe_to_blender.send(Result(repo.to_dict()))

def debug_hang():
    """Hangs for an hour. For testing purposes only."""

    import time
    time.sleep(3600)