­­­­­­­­­­­­­­­­­­
<!DOCTYPE html>
<html>
"""
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 re
from typing import Generator, Union

from defence360agent.contracts.messages import MessageType
from defence360agent.contracts.plugins import MessageSink, expect
from defence360agent.model.simplification import run_in_executor
from defence360agent.utils import is_cluster, nice_iterator
from imav.malwarelib.config import MalwareScanType, VulnerabilityHitStatus
from imav.contracts.plugins import ProcessOrder
from imav.malwarelib.model import MalwareScan, VulnerabilityHit
from imav.malwarelib.tenant_path import (
    TenantPath,
    split_prefixed,
    to_prefixed,
)


class StoreVulnerabilities(MessageSink):
    PROCESSING_ORDER = ProcessOrder.AFTER_STORE_SCAN

    async def create_sink(self, loop):
        self._loop = loop

    @staticmethod
    def get_outdated_entries(
        path_obj: Union[str, list],
        scan_type: str | None = None,
    ) -> Generator[str, None, None]:
        """
        Return files that may already not be vulnerable, yet we still
        consider them such.

        For example, a vulnerable file might have been removed manually.

        Paths arrive in the stored form (tenant-prefixed on k8s), so the
        matching is inherently scoped to the owning tenant.
        """
        # NOTE: this logic was taken by analogy with StoreMalwareHits
        # consider optimizing this code
        paths = [path_obj] if isinstance(path_obj, str) else path_obj
        if scan_type == MalwareScanType.REALTIME:
            # already the stored form; delete_hits matches it directly
            yield from paths
            return
        for target_path in paths:
            # the tenant rides the prefix; the glob side needs the bare
            # in-container remainder
            tenant, bare = (
                split_prefixed(target_path)
                if is_cluster()
                else (None, target_path)
            )
            for path in TenantPath.iglob(bare, user=tenant):
                path = str(path.resolve())
                # rows are stored tenant-prefixed on k8s; compare and yield
                # in that form so delete_hits matches
                db_path = to_prefixed(path, tenant)
                file_clauses = (
                    VulnerabilityHit.orig_file == db_path
                ) & VulnerabilityHit.status.in_(
                    [
                        VulnerabilityHitStatus.VULNERABLE,
                        VulnerabilityHitStatus.REVERTED,
                    ]
                )
                if (
                    TenantPath(path, user=tenant).is_file()
                    and VulnerabilityHit.select().where(file_clauses).first()
                ):
                    yield db_path
                else:
                    scanned_dir = re.escape(db_path) + r"(/.*|\b)"
                    dir_clauses = VulnerabilityHit.orig_file.regexp(
                        scanned_dir
                    ) & VulnerabilityHit.status.in_(
                        [
                            VulnerabilityHitStatus.VULNERABLE,
                            VulnerabilityHitStatus.REVERTED,
                        ]
                    )
                    yield from (
                        i.orig_file
                        for i in VulnerabilityHit.select().where(dir_clauses)
                    )

    def _delete_outdated_entries(self, summary: dict) -> None:
        file_patterns = summary.pop("file_patterns", None)
        exclude_patterns = summary.pop("exclude_patterns", None)
        if (
            summary.get("error") is None
            and file_patterns is None
            and exclude_patterns is None
        ):
            outdated_entries = self.get_outdated_entries(
                summary["path"],
                scan_type=summary["type"],
            )
            VulnerabilityHit.delete_hits(outdated_entries)

    @expect(MessageType.MalwareScan)
    async def process_hits(self, message):
        if not message["summary"].get("started") or message["results"] is None:
            # Scan is queued/aborted.
            return

        scan = MalwareScan.get(scanid=message["summary"]["scanid"])
        # get('path') indicates that this is the second message,
        # even if they are out of order
        if message["summary"].get("path") is not None:
            # keep the same logic as for malware hits: the sweep may block
            # on syncer stat/glob round trips in cluster mode
            await run_in_executor(
                self._loop, self._delete_outdated_entries, message["summary"]
            )
        if results := {
            filename: data
            for filename, data in message["results"].items()
            if VulnerabilityHit.match(data["hits"][0]["matches"])
        }:
            # TODO: handle possible races when we implement patch/revert
            async for filename, data in nice_iterator(results.items()):
                # the latest detection type is relevant
                VulnerabilityHit.create(
                    scanid=scan.scanid,
                    owner=data["owner"],
                    user=data["user"],
                    size=data["size"],
                    hash=data["hash"],
                    # DB stores the canonical tenant-prefixed form on k8s
                    orig_file=to_prefixed(filename, data["user"]),
                    type=data["hits"][0]["matches"],
                    timestamp=data["hits"][0]["timestamp"],
                    status=VulnerabilityHitStatus.VULNERABLE,
                )