OwlCyberSecurity - MANAGER

Edit File: vulnerabilities.py

"""
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/>.

This software is also available under ImunifyAV commercial license,
see <https://www.imunify360.com/legal/eula>
"""
from defence360agent.contracts.messages import MessageType
from defence360agent.rpc_tools.lookup import (
    CommonEndpoints,
    bind,
)
from imav.malwarelib.config import VulnerabilityHitStatus
from imav.malwarelib.api.vulnerability import VulnerabilityAPI
from imav.malwarelib.model import VulnerabilityHit
from imav.malwarelib.vulnerabilities.storage import restore_hits

class VulnerabilitiesEndpoints(CommonEndpoints):
    async def _get_vulnerabilities_details(self, hits: list):
        vuln_ids = set()
        for hit in hits:
            vuln_ids |= set(
                VulnerabilityHit.get_vulnerability_ids(hit["type"])
            )
        return await VulnerabilityAPI.get_details(vuln_ids)

@bind("vulnerabilities", "file", "list")
    async def vulnerabilities_file_list(self, user=None, **kwargs):
        """
        Return list vulnerable/patched files
        """
        max_count, hits = VulnerabilityHit.list(user=user, **kwargs)
        vuln_info = await self._get_vulnerabilities_details(hits)
        results = []
        for hit in hits:
            record = {
                "id": hit["id"],
                "username": hit["username"],
                "file_path": hit["file_path"],
                "status": hit["status"],
                "app_name": "",
                "vulnerabilities": [],
            }
            for vuln_id in VulnerabilityHit.get_vulnerability_ids(hit["type"]):
                record["vulnerabilities"].append(
                    {
                        "cve_id": vuln_info[vuln_id]["cveId"],
                        "vulnerability_type": vuln_info[vuln_id]["type"],
                        "vulnerability_description": vuln_info[vuln_id][
                            "name"
                        ],
                    }
                )
                if not record["app_name"]:  # set it once
                    record["app_name"] = vuln_info[vuln_id]["app"]
            results.append(record)
        return max_count, results

@bind("vulnerabilities", "file", "patch")
    async def vulnerabilities_file_patch(self, paths, user=None):
        await self._sink.process_message(
            MessageType.VulnerabilityPatchTask(filelist=paths, initiator=user)
        )

@bind("vulnerabilities", "file", "revert")
    async def vulnerabilities_file_revert(self, paths, user=None):
        query = VulnerabilityHit.select().where(
            VulnerabilityHit.orig_file.in_(paths),
            VulnerabilityHit.status.in_([VulnerabilityHitStatus.PATCHED]),
        )
        if user is not None:
            query.where(VulnerabilityHit.user.in_([user]))
        hits = list(query)
        succeded, failed = await restore_hits(hits)
        return {
            "succeded": [hit.orig_file for hit in succeded],
            "failed": [hit.orig_file for hit in failed],
        }