Mini Shell

import asyncio
import functools
import logging
import os
from tempfile import TemporaryDirectory

from . import helpers
from .backup_backends import plesk
from .backup_backends_lib import FtpBackupBase
from .config import Flags
from .safe_fileops import safe_move
from .scan import scan

RESTORE_DIR_NAME = '.restore-infected'
RESTORE_DIR = os.path.expanduser(os.path.join('~', RESTORE_DIR_NAME))

logger = logging.getLogger(__name__)


class FileData:
    def __init__(self, filename):
        self.filename = os.path.abspath(filename)
        if os.path.exists(self.filename):
            stat = os.stat(self.filename)
            self.size = stat.st_size
            self.mtime = helpers.DateTime.fromtimestamp(stat.st_mtime)
        else:
            self.size = None
            self.mtime = None

    def __eq__(self, other):
        return (self.size, self.mtime, self.filename) == (
            other.size,
            other.mtime,
            other.filename,
        )

    def __repr__(self):
        return '<{0}({1})>'.format(self.__class__.__name__, self.filename)


class InfectedList:
    def __init__(self, files, restore_dir, scan_func):
        self.files = [FileData(f) for f in files]
        self.restore_dir = restore_dir
        self.scan_func = scan_func

    def __bool__(self):
        return bool(self.files)

    def __iter__(self):
        return iter(self.files)

    def check(self, backup, check_list, **kw):
        """
        Return list of cured files. This files are removed from InfectedList
        """
        check_map = backup.restore(check_list, self.restore_dir, **kw)
        scan_list = check_map.keys()

        still_infected = self.scan_func(scan_list)

        return self._build_cured_list(check_map, scan_list, still_infected)

    async def async_check(self, backup, check_list, **kw):
        """
        Return list of cured files. This files are removed from InfectedList
        """
        loop = asyncio.get_event_loop()
        check_map = await loop.run_in_executor(
            None,
            functools.partial(
                backup.restore, check_list, self.restore_dir, **kw
            ),
        )
        scan_list = check_map.keys()

        still_infected = await self.scan_func(scan_list)

        return self._build_cured_list(check_map, scan_list, still_infected)

    def _build_cured_list(self, check_map, scan_list, still_infected):
        cured = set(scan_list) - set(still_infected)
        cured_list = []
        for src in cured:
            dst = check_map[src]
            try:
                self.files.remove(FileData(dst))
            except ValueError:
                logger.warning(
                    'Skipping restoration of a changed file %s', dst
                )
                self.files = [f for f in self.files if f.filename != dst]
            else:
                cured_list.append((src, dst))
        return cured_list


def prep_restore_dir(tmp_dir=None):
    restore_dir = RESTORE_DIR
    try:
        if tmp_dir:
            restore_dir = os.path.join(tmp_dir, RESTORE_DIR_NAME)
        os.makedirs(restore_dir, exist_ok=True)
    except FileExistsError as e:
        message = os.linesep.join([
            str(e),
            '',
            'There is a file in place of {}'.format(restore_dir),
            'Most probably it is done to prevent `restore` to be run',
            'Remove this file to be able to restore again'
        ])
        raise FileExistsError(message) from e
    else:
        return TemporaryDirectory(dir=restore_dir)


def restore_infected(
    backend,
    files,
    until=None,
    scan_func=scan,
    pre_restore_hook=...,
    tmp_dir=None,
    **kw,
):
    restore_dir = prep_restore_dir(tmp_dir=tmp_dir)

    infected_list = InfectedList(files, restore_dir.name, scan_func)

    args = backend.pre_backups(files, until)
    args = args or {}

    success = []
    failed = []
    for backup in backend.backups(
            until=until, **args, async_=False, tmp_dir=tmp_dir
    ):
        if Flags.DisableFtpBackups and isinstance(backup, FtpBackupBase):
            continue
        check_list = []
        for infected_file in infected_list:
            try:
                backup_file = backup.file_data(infected_file.filename)
            except FileNotFoundError:
                continue
            except EOFError:
                logger.warning(
                    "Unexpected end of file reached during processing backup: "
                    "{}. Skipping...".format(backup)
                )
                if backend is plesk:
                    logger.warning(
                        "Plesk multivolume backups currently are not supported"
                    )
                continue
            if infected_file == backup_file:
                continue
            check_list.append(backup_file)

        if check_list:
            restore_list = infected_list.check(backup, check_list, **kw)
            for src, dst in restore_list:
                try:
                    safe_move(src, dst)
                except PermissionError:
                    failed.append(dst)
                else:
                    success.append(dst)

        backup.close()
        if not infected_list:
            break

    failed.extend([f.filename for f in infected_list])

    backend.cleanup()
    return success, failed


async def async_restore_infected(
    backend,
    files,
    until=None,
    scan_func=scan,
    pre_restore_hook=...,
    tmp_dir=None,
    **kw,
):
    restore_dir = prep_restore_dir(tmp_dir=tmp_dir)

    infected_list = InfectedList(files, restore_dir.name, scan_func)

    args = await backend.pre_backups(files, until, async_=True)
    args = args or {}
    success = []
    failed = []
    loop = asyncio.get_event_loop()
    for backup in await backend.backups(
            until=until, **args, async_=True, tmp_dir=tmp_dir
    ):
        if Flags.DisableFtpBackups and isinstance(backup, FtpBackupBase):
            continue
        check_list = []
        for infected_file in infected_list:
            try:
                backup_file = await loop.run_in_executor(
                    None,
                    functools.partial(
                        backup.file_data, infected_file.filename
                    ),
                )
            except FileNotFoundError:
                continue
            if infected_file == backup_file:
                continue
            check_list.append(backup_file)

        if check_list:
            restore_list = await infected_list.async_check(
                backup, check_list, **kw
            )

            for src, dst in restore_list:
                try:
                    safe_move(src, dst)
                except PermissionError:
                    failed.append(dst)
                else:
                    success.append(dst)

        backup.close()
        if not infected_list:
            break

    failed.extend([f.filename for f in infected_list])

    await backend.cleanup()
    return success, failed