Edit

[EDIT] FILE: patcher.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>
"""
import asyncio
import logging
import os
import subprocess
import tempfile
import time
from contextlib import suppress
from typing import Dict, List, Optional, Tuple, Union

from defence360agent.contracts.config import MalwareSignatures
from defence360agent.contracts.messages import MessageType
from defence360agent.utils import base64_encode_filename
from imav.contracts.config import MalwareTune
from imav.malwarelib.engine import curator_path
from imav.malwarelib.model import VulnerabilityHit
from imav.malwarelib.utils.revisium import (
    DeletionType,
    ErrorType,
    RescanResultType,
    RevisiumCSVFile,
    RevisiumJsonFile,
    RevisiumTempFile,
)

logger = logging.getLogger(__name__)

def patcher_result_instance(tempdir=None, mode=None):
    if MalwareTune.USE_JSON_REPORT:
        return RevisiumJsonFile(tempdir, mode)
    return RevisiumCSVFile(tempdir, mode)

class VulnerabilityPatcherLog(RevisiumTempFile):
    pass

class VulnerabilityPatcherProgress(RevisiumJsonFile):
    pass

class VulnerabilityPatchFileList(RevisiumTempFile):
    def write(self, filelist):
        with self._path.open("wb") as w:
            w.writelines(base64_encode_filename(f) + b"\n" for f in filelist)

def _parse_int(value: Union[str, int]) -> int:
    """Convert str|int to int, in case errors return -2
    -1 used as default value when storing CH
    """
    try:
        return int(value)
    except ValueError:
        return -2

class PatchResultEntry(dict):
    def __init__(self, data: Dict[str, Union[str, int]]):
        # fields:
        # d - patch result
        # e - error description
        # s - signature that was triggered for the file during scan
        # f - file path (it's unexpected that f is absent)
        # r - the result of aibolit rescan after patch
        #
        # We shouldn't fail on parsing one record (to do not stop processing
        #  report), so we consider default values for all fields.
        super().__init__(
            d=_parse_int(data.get("d", -1)),
            e=_parse_int(data.get("e", -1)),
            s=data["s"],
            f=data["f"],
            r=_parse_int(data.get("r", -1)),
            mtime_before=_parse_int(data.get("mb", -1)),
            mtime_after=_parse_int(data.get("ma", -1)),
            hash_before=data.get("hb", ""),
            hash_after=data.get("ha", ""),
        )

def is_patched(self):
        if (
            self.is_failed() or self.requires_myimunify_protection()
        ):  # is myimunify_protection relevant for the patch?
            return False

if self["e"] == ErrorType.NOT_CLEANEDUP:
            logger.warning(
                "File has changed, assuming that it was patched: %s", self["f"]
            )
            return True

return (
            self["e"] == ErrorType.NO_ERROR
            and self["d"] == DeletionType.PATCH_APPLIED
        )

def is_failed(self):
        return (
            self["r"] == RescanResultType.DETECTED
        )  # TODO: check if it is relevant for the patch

def requires_myimunify_protection(self):
        return self["r"] == RescanResultType.REQUIRED_ADVANCED_SIGNATURES

def not_exist(self):
        return not self.is_failed() and self["e"] == ErrorType.FILE_NOT_EXISTS

class PatchResult(Dict[str, PatchResultEntry]):
    """
    Cleanup result container for result entries
    """

def __init__(self, report=None):
        if report:
            super().__init__({e["f"]: PatchResultEntry(e) for e in report})

@staticmethod
    def __key(hit: Union[str, VulnerabilityHit]) -> str:
        return getattr(hit, "orig_file", hit)

def __contains__(self, hit: Union[str, VulnerabilityHit]):
        return super().__contains__(self.__key(hit))

def __getitem__(self, hit: Union[str, VulnerabilityHit]):
        return super().__getitem__(self.__key(hit))

class VulnerabilityPatcher:
    PROCU_DB = MalwareSignatures.PROCU_DB

def __init__(self, loop=None, sink=None):
        self._loop = loop if loop else asyncio.get_event_loop()
        self._sink = sink

def _cmd(
        self,
        filename,
        progress_path,
        result_path,
        log_path,
        *,
        username,
        use_csv=True,
    ):
        patcher_binary = curator_path()
        cmd = [
            "/opt/ai-bolit/wrapper",
            patcher_binary,
            "--deobfuscate",
            "--nobackup",
            "--patch-vulnerabilities",
            "--rescan",
            "--list=%s" % filename,
            "--input-fn-b64-encoded",
            "--username=%s" % username,
            "--report-hashes",
            "--log=%s" % log_path,
            "--progress=%s" % progress_path,
        ]
        if use_csv:  # TODO: is this still needed?
            cmd.extend(["--csv_result=%s" % result_path])
        else:
            cmd.extend(["--result=%s" % result_path])

if os.path.exists(self.PROCU_DB):
            cmd.append("--avdb")
            cmd.append(self.PROCU_DB)

return cmd

@staticmethod
    def _get_patcher_error_info(
        exc: Exception,
        cmd: List[str],
        returncode: int,
        stdout: Optional[bytes],
        stderr: Optional[bytes],
    ):
        return dict(
            exception=exc.__class__.__name__,
            return_code=returncode,
            command=cmd,
            out=stdout.decode(errors="replace") if stdout is not None else "",
            err=stderr.decode(errors="replace") if stderr is not None else "",
        )

async def _send_failed_message(self, info: dict):
        if self._sink:
            try:
                msg = MessageType.VulnerabilityPatchFailed(
                    {**info, **{"timestamp": int(time.time())}}
                )
                await self._sink.process_message(msg)
            except asyncio.CancelledError:
                raise
            except Exception:
                logger.exception("Exception while sending PatchFailed message")

async def start(
        self,
        user,
        filelist,
    ) -> Tuple[PatchResult, Optional[str], List[str]]:
        tempdir = tempfile.gettempdir()
        result_file = patcher_result_instance(tempdir=tempdir)
        use_csv = isinstance(result_file, RevisiumCSVFile)

with VulnerabilityPatchFileList(
            tempdir=tempdir, mode=0o644
        ) as flist, VulnerabilityPatchFileList(
            tempdir=tempdir, mode=0o644
        ), VulnerabilityPatcherProgress(
            tempdir=tempdir
        ) as progress, result_file as result, VulnerabilityPatcherLog(
            tempdir=tempdir
        ) as log:
            flist.write(filelist)
            cmd = self._cmd(
                flist.filename,
                progress.filename,
                result.filename,
                log.filename,
                username=user,
                use_csv=use_csv,
            )
            logger.debug("Executing %s", " ".join(cmd))

out, err = b"", b""
            proc = None
            try:
                proc = await asyncio.subprocess.create_subprocess_exec(
                    *cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                )
                out, err = await proc.communicate()
                report = result.read()
            except asyncio.CancelledError:
                if proc:
                    with suppress(ProcessLookupError):
                        proc.terminate()
                raise
            except Exception as exc:
                info = self._get_patcher_error_info(
                    exc,
                    cmd,
                    proc.returncode if proc else 126,  # 126 - permission error
                    stdout=out,
                    stderr=err,
                )
                # Group errors by exit code on Sentry
                logger.error(
                    "Patch vulnerabilities failed"
                    f" exit_code={info.get('return_code')}: %s",
                    f"{info.get('out')} {info.get('err')}",
                    extra={**info, "exception": exc},
                )
                await self._send_failed_message(
                    {**info, **dict(message=str(exc))}
                )
                return PatchResult(), repr(exc), cmd

return PatchResult(report), None, cmd

CANCEL