malware_transform.py

# (C) Copyright IBM Corp. 2024.
# Licensed under the Apache License, Version 2.0 (the “License”);
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#  http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an “AS IS” BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
################################################################################
import logging
import subprocess
import time
from argparse import ArgumentParser, Namespace
from io import BytesIO
from typing import Any

import pyarrow as pa
from data_processing.transform import AbstractTableTransform, TransformConfiguration
from data_processing.utils.transform_utils import TransformUtils


INPUT_COLUMN_KEY = "malware_input_column"
OUTPUT_COLUMN_KEY = "malware_output_column"
DEFAULT_INPUT_COLUMN = "contents"
DEFAULT_OUTPUT_COLUMN = "virus_detection"
CLAMD_TIMEOUT_SEC = 180
INIT_TIMEOUT_SEC = 60


# def check_clamd():
#     logger = get_logger(__name__)
#     cd = clamd.ClamdUnixSocket()
#     check_end = time.time() + INIT_TIMEOUT_SEC
#     while True:
#         try:
#             cd.ping()
#             break
#         except Exception as err:
#             if time.time() < check_end:
#                 logger.debug("waiting until clamd is ready...")
#                 time.sleep(1)
#             else:
#                 logger.error(
#                     f"clamd didn't become ready in {INIT_TIMEOUT_SEC} sec. Please check if clamav container is running."
#                 )
#                 raise err
#     del cd


def check_clamd(logger: logging.Logger):
    import clamd

    cd = clamd.ClamdUnixSocket()
    check_end = time.time() + INIT_TIMEOUT_SEC
    while True:
        try:
            cd.ping()
            break
        except Exception as err:
            if time.time() < check_end:
                logger.debug("waiting until clamd is ready...")
                time.sleep(1)
            else:
                logger.error(
                    f"clamd didn't become ready in {INIT_TIMEOUT_SEC} sec. Please check if clamav container is running."
                )
                raise err
    del cd


class MalwareTransform(AbstractTableTransform):
    """
    Implements a simple copy of a pyarrow Table.
    """

    def __init__(self, config: dict[str, Any]):
        """
        Initialize based on the dictionary of configuration information.
        This is generally called with configuration parsed from the CLI arguments defined
        by the companion runtime, MalwareTransformRuntime.  If running inside the RayMutatingDriver,
        these will be provided by that class with help from the RayMutatingDriver.
        """
        import clamd

        # Make sure that the param name corresponds to the name used in apply_input_params method
        # of MalwareTransformConfiguration class
        super().__init__(config)
        self.warning_issued = False
        self.input_column = config.get(INPUT_COLUMN_KEY, DEFAULT_INPUT_COLUMN)
        self.output_column = config.get(OUTPUT_COLUMN_KEY, DEFAULT_OUTPUT_COLUMN)
        # Check local clamd process is running
        cd = clamd.ClamdUnixSocket()
        try:
            cd.ping()
        except clamd.ConnectionError:
            self.logger.info("Clamd process is not running. Start clamd process.")
            subprocess.Popen("clamd", shell=True)
            # Wait for starting up clamd
            timeout = time.time() + CLAMD_TIMEOUT_SEC
            while True:
                try:
                    cd.ping()
                    break
                except clamd.ConnectionError as err:
                    if time.time() < timeout:
                        self.logger.debug("Clamd is not ready. Retry after 5 seconds.")
                        time.sleep(5)
                    else:
                        self.logger.error(f"Clamd didn't become ready in {CLAMD_TIMEOUT_SEC} seconds.")
                        raise err
        del cd

    def transform(self, table: pa.Table, file_name: str = None) -> tuple[list[pa.Table], dict[str, Any]]:
        """
        Put Transform-specific to convert one Table to 0 or more tables. It also returns
        a dictionary of execution statistics - arbitrary dictionary
        This implementation makes no modifications so effectively implements a copy of the
        input parquet to the output folder, without modification.
        """
        import clamd

        self.logger.debug(f"Transforming one table with {len(table)} rows")

        cd = clamd.ClamdUnixSocket()

        def _scan(content: str) -> str | None:
            if content is None:
                return None
            (status, description) = cd.instream(BytesIO(content.encode()))["stream"]
            if status == "FOUND":
                self.logger.debug(f"Detected: {description}")
                return description or "UNKNOWN"
            return None

        virus_detection = pa.array(list(map(_scan, table[self.input_column].to_pylist())), type=pa.string())

        nrows = table.num_rows
        clean = virus_detection.null_count
        infected = nrows - clean
        table = TransformUtils.add_column(table, self.output_column, virus_detection)
        # Add some sample metadata.
        self.logger.debug(f"Virus detection {infected} / {nrows} rows")
        metadata = {"clean": clean, "infected": infected}
        return [table], metadata


class MalwareTransformConfiguration(TransformConfiguration):

    """
    Provides support for configuring and using the associated Transform class include
    configuration with CLI args and combining of metadata.
    """

    def __init__(self):
        super().__init__(
            name="Malware",
            transform_class=MalwareTransform,
        )
        from data_processing.utils import get_logger

        self.logger = get_logger(__name__)

    def add_input_params(self, parser: ArgumentParser) -> None:
        """
        Add Transform-specific arguments to the given  parser.
        This will be included in a dictionary used to initialize the MalwareTransform.
        By convention a common prefix should be used for all transform-specific CLI args
        (e.g, noop_, pii_, etc.)
        """
        parser.add_argument(
            f"--{INPUT_COLUMN_KEY}",
            type=str,
            default=DEFAULT_INPUT_COLUMN,
            help="input column name",
        )
        parser.add_argument(
            f"--{OUTPUT_COLUMN_KEY}",
            type=str,
            default=DEFAULT_OUTPUT_COLUMN,
            help="output column name",
        )

    def apply_input_params(self, args: Namespace) -> bool:
        """
        Validate and apply the arguments that have been parsed
        :param args: user defined arguments.
        :return: True, if validate pass or False otherwise
        """
        self.params[INPUT_COLUMN_KEY] = args.malware_input_column
        self.params[OUTPUT_COLUMN_KEY] = args.malware_output_column
        self.logger.info(f"malware parameters are : {self.params}")
        return True