"""
This program is free software: you can redistribute it and/or modify it under
the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License,
or (at your option) any later version.


This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. 
See the GNU General Public License for more details.


You should have received a copy of the GNU General Public License
 along with this program.  If not, see <https://www.gnu.org/licenses/>.

Copyright © 2019 Cloud Linux Software Inc.

This software is also available under ImunifyAV commercial license,
see <https://www.imunify360.com/legal/eula>
"""
import logging
import pwd
import asyncio

from defence360agent.contracts.hook_events import HookEvent
from defence360agent.contracts.config import (
    choose_value_from_config,
    SystemConfig,
)
from defence360agent.contracts.messages import MessageType
from defence360agent.contracts.config import ANTIVIRUS_MODE
from defence360agent.contracts.plugins import (
    MessageSink,
    MessageSource,
    expect,
)
from defence360agent.subsys.panels.hosting_panel import HostingPanel
from defence360agent.subsys.persistent_state import load_state, save_state
from defence360agent.utils import Scope

if not ANTIVIRUS_MODE:
    # imav.wordpress.plugin imports clcommon which require pymysql in Plesk
    # environment. Which is available only im360 environment. (venv-firewall)
    # Before moving it to venv package we need to make sure that it is not
    # increasing memory usage.
    from imav.wordpress import plugin
else:
    plugin = None
from imav.contracts.config import Wordpress
from imav.malwarelib.model import MalwareHit
from imav.model.wordpress import WPSite

logger = logging.getLogger(__name__)


class ImunifySecurityPlugin(MessageSink, MessageSource):
    SCOPE = Scope.IM360

    def __init__(self):
        self._loop = None
        self._sink = None
        self.installation_completed = load_state("ImunifySecurityPlugin").get(
            "installed"
        )
        self.last_config_value = (
            load_state("ImunifySecurityPlugin").get("enabled")
            or Wordpress.SECURITY_PLUGIN_ENABLED
        )
        self.installation_task: asyncio.Task | None = None
        self.deleting_task: asyncio.Task | None = None

    async def create_sink(self, loop):
        pass

    async def create_source(self, loop, sink):
        self._loop = loop
        self._sink = sink

    def _task_in_progress(self, task_attr_name):
        if not hasattr(self, task_attr_name):
            logger.error("Unknown task '%s'", task_attr_name)
        task = getattr(self, task_attr_name)

        return task is not None and not task.done() and not task.cancelled()

    async def process_installation(self, coro, for_new_sites=False):
        if self._task_in_progress("deleting_task"):
            if for_new_sites:
                return

            self.deleting_task.cancel()
            await self.deleting_task

        if self._task_in_progress("installation_task"):
            logger.warning("Installation is already running")
            return

        self.installation_task = asyncio.create_task(coro)

    async def process_deleting(self, coro):
        if self._task_in_progress("installation_task"):
            self.installation_task.cancel()
            await self.installation_task

        if self._task_in_progress("deleting_task"):
            logger.warning("Deleting is already running")
            return

        self.deleting_task = asyncio.create_task(coro)

    async def get_valid_users(self):
        users = await HostingPanel().get_users()
        valid_users = set()
        invalid_users = set()

        for user in users:
            default_action, _ = choose_value_from_config(
                "MALWARE_SCANNING",
                "default_action",
                username=user,
            )

            if default_action == "cleanup":
                valid_users.add(user)
            else:
                invalid_users.add(user)

        if invalid_users:
            logger.info(
                "WordPress plugin installation is currently available only"
                ' with the following config settings: {"MALWARE_SCANNING":'
                ' {"default_action": "cleanup"}.  Additional options will be'
                " supported later."
            )
            logger.info(f"Skipping installation for {invalid_users=}")
        return valid_users

    @expect(MessageType.WordpressPluginAction)
    async def manage_plugin_action(self, message):
        logger.info(
            "ImunifySecurityPlugin received message action: %s method: %s",
            message.action,
            message.method,
        )

        if message.action == "install_on_new_sites":
            if not self.installation_completed:
                # The installation is not completed yet. We cannot know reliably which sites are new.
                return

            valid_users = await self.get_valid_users()
            await self.process_installation(
                plugin.install_for_users(users=valid_users, sink=self._sink),
                for_new_sites=True,
            )
        elif message.action == "tidy_up":
            # Tidy up sites form which the WordPress plugin was deleted manually.
            await plugin.tidy_up_manually_deleted(sink=self._sink)

            if not Wordpress.SECURITY_PLUGIN_ENABLED:
                # If the feature is disabled, remove the WordPress plugin from all sites.
                await self.process_deleting(
                    plugin.remove_all_installed(sink=self._sink)
                )
                self.installation_completed = False

    @expect(MessageType.ConfigUpdate)
    async def manage_plugin_installation(self, message):
        if not isinstance(message["conf"], SystemConfig):
            return

        current_config_value = Wordpress.SECURITY_PLUGIN_ENABLED
        if current_config_value == self.last_config_value:
            return

        # Update last config value immediately to prevent multiple installations
        self.last_config_value = current_config_value

        if current_config_value and not self.installation_completed:
            valid_users = await self.get_valid_users()

            await self.process_installation(
                plugin.install_for_users(users=valid_users, sink=self._sink)
            )

            self.installation_completed = True

        elif not current_config_value and self.installation_completed:
            await self.process_deleting(
                plugin.remove_all_installed(sink=self._sink)
            )
            self.installation_completed = False

        save_state(
            "ImunifySecurityPlugin",
            {
                "installed": self.installation_completed,
                "enabled": current_config_value,
            },
        )

    @expect(HookEvent.MalwareCleanupFinished)
    async def handle_malware_cleanup_finished(self, message):
        """
        INFO    [2025-02-24 12:00:20,384] imav.plugins.wordpress: Malware cleanup finished:
        HookEvent.MalwareCleanupFinished(
            {
                'cleanup_id': 'fa4fe7e48dbf45588f53b24366cd8893',
                'started': 1740398411.786418,
                'error': None,
                'total_files': 3,
                'total_cleaned': 3,
                'status': 'ok'
            }
        )
        """
        # Skip if plugin is disabled
        if not self.last_config_value:
            return

        # Leave early if status is not ok or the started time is missing.
        if message.get("status") != "ok" or not message.get("started"):
            return

        # load all malware hits cleaned since the cleanup started
        hits = MalwareHit.cleaned_since(message["started"])

        site_paths = set()

        # Collect all site paths that need to be updated.
        for hit in hits:
            if hit.resource_type == "file":
                try:
                    user_info = pwd.getpwnam(hit.user)
                    uid = user_info.pw_uid
                    user_sites = plugin.get_sites_for_user(uid)

                    for site_path in user_sites:
                        if hit.orig_file.startswith(site_path):
                            site_paths.add((site_path, uid))
                            break

                except KeyError:
                    pass

        if not site_paths:
            logger.debug("Cleanup finished => no sites found for cleaned hits")
            return

        logger.info(
            "Cleanup finished => %s site(s) need to be updated",
            len(site_paths),
        )

        # Convert paths to WPSite objects with empty domain and update data on the sites that need to be updated.
        # We need to work with paths here because sometimes the domain is not set, see https://cloudlinux.atlassian.net/browse/DEF-32238.
        wordpress_sites = [
            WPSite(docroot=site_path, domain="", uid=uid)
            for site_path, uid in site_paths
        ]

        await plugin.update_data_on_sites(self._sink, wordpress_sites)

        logger.info("%s site(s) updated after a cleanup", len(wordpress_sites))

    @expect(HookEvent.MalwareScanningFinished)
    async def handle_malware_scan_finished(self, message):
        """
        INFO    [2025-02-24 11:57:17,968] imav.plugins.wordpress: Malware scan finished:
        HookEvent.MalwareScanningFinished(
            {
                'scan_id': 'b9bd136aff0a4d87a248c859cfe41c47',
                'scan_type': 'user',
                'path': '/home/user1'
            }
        )
        INFO    [2025-02-24 12:00:10,740] imav.plugins.wordpress: Malware scan finished:
        HookEvent.MalwareScanningFinished(
            {
                'scan_id': 'a74271d2cdd04e0c9bd49ef6de23e0d8',
                'scan_type': 'user',
                'path': '/home/user4',
                'started': 1740398383,
                'total_files': 39229,
                'total_malicious': 3,
                'error': None,
                'status': 'ok',
                'scan_params': {'intensity_cpu': 2, 'intensity_io': 2, 'intensity_ram': 2048, 'initiator': None, 'file_patterns': None, 'exclude_patterns': None, 'follow_symlinks': False, 'detect_elf': True},
                'stats': {'scan_time': 27, 'mem_peak': 28217344, 'smart_time_hs': 0.004, 'scan_time_hs': 1.1751, 'smart_time_preg': 0, 'scan_time_preg': 2.7391, 'finder_time': 13.5896, 'cas_time': 0.7562, 'deobfuscate_time': 0.8998, 'total_files': 39229}
            }
        )
        """
        # Skip if plugin is disabled
        if not self.last_config_value:
            return

        # Leave early if status is not ok or path or stats are missing.
        if (
            message.get("status") != "ok"
            or not message.get("path")
            or not message.get("stats")
        ):
            return

        # Malware scan is finished, figure out what sites need to be updated based on the path.
        path = message["path"]
        sites = plugin.get_sites_by_path(path)
        if not sites:
            logger.debug("Scan finished => no sites found for path=%s", path)
            return

        # Update data on the sites that need to be updated.
        logger.info(
            "Scan finished => %s site(s) need to be updated", len(sites)
        )

        await plugin.update_data_on_sites(self._sink, sites)

        logger.info("%s site(s) updated after a scan", len(sites))