EVOLUTION-NINJA

Edit File: aibolit_job.py

#!/usr/bin/env python3
"""Helper functions to kick off the aibolit scan on file upload.

See {modsec,pureftpd}-on-upload scripts for the example usage.

"""
import base64
import json
import os
import socket
import time
import uuid
from contextlib import ExitStack, suppress
from pathlib import Path
from tempfile import NamedTemporaryFile

# NOTE: despite being placed inside im360, the aibolit_job package
# shouldn't import anything from the agent

__all__ = [
    "RESIDENT_DIR",
    "RESIDENT_IN_DIR_NOTIFY_REL_PATH",
    "RESIDENT_IN_DIR_UPLOAD_REL_PATH",
    "UPLOAD_TIMEOUT",
    "create_notify_job",
    "create_remaining_time_func",
    "create_upload_job",
    "notify_aibolit_start_it_if_necessary",
]

AIBOLIT_STARTUP_SOCKET = "/var/run/defence360agent/aibolit-resident.sock"
RESIDENT_DIR = Path("/var/imunify360/aibolit/resident")
# path for uploads jobs relative to RESIDENT_DIR
RESIDENT_IN_DIR_UPLOAD_REL_PATH = Path("in/upload-jobs")
# path for notify jobs relative to RESIDENT_DIR
RESIDENT_IN_DIR_NOTIFY_REL_PATH = Path("in/notify-jobs")
STUCK_TIMEOUT = 5
UPLOAD_TIMEOUT = 10  # file lock timeout

def create_remaining_time_func(
    timeout, *, start_time=None, timer=time.monotonic
):
    """Create remaining_time() function.

Start the timer if start_time is None otherwise use given value as
    the start time.

remaining_time() raises TimeoutError in *timeout*
    seconds after the *start_time* according to the *timer*.

"""
    start_time = timer()

def remaining_time():
        time_left = timeout - (timer() - start_time)
        if time_left <= 0:
            raise TimeoutError
        return time_left

return remaining_time

def create_upload_job(files, *, resident_dir_path=RESIDENT_DIR, timeout=None):
    """Create PID.upload_job in the resident/in dir."""
    create_job(
        files,
        job_path=resident_dir_path
        / RESIDENT_IN_DIR_UPLOAD_REL_PATH
        / "{}.upload_job".format(os.getpid()),
        timeout=timeout,
    )

def create_notify_job(files, *, resident_dir_path=RESIDENT_DIR, timeout=None):
    """Create UUID.notify_job in the resident/in dir."""
    create_job(
        files,
        job_path=resident_dir_path
        / RESIDENT_IN_DIR_NOTIFY_REL_PATH
        / "{}.notify_job".format(uuid.uuid4()),
        timeout=timeout,
    )

def create_job(files, *, job_path, timeout=None):
    # note: timeout is ignored to avoid complicating the common case
    # when the file operations take less than the timeout
    files = [base64.b64encode(os.fsencode(fn)).decode("ascii") for fn in files]
    write_atomic(
        job_path,
        json.dumps({"files": files}).encode("ascii"),
    )

def notify_aibolit_start_it_if_necessary():
    """Notify aibolit that there is a file to scan.

Start/restart aibolit if necessary.
    """
    # The wakeup byte socket-activates aibolit-resident on cold start;
    # once running it polls the job dir, so a dropped byte is harmless.
    with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as sock:
        sock.setblocking(False)
        sock.connect(AIBOLIT_STARTUP_SOCKET)
        try:
            sock.send(b"1")
        except (BlockingIOError, InterruptedError):
            pass

def write_atomic(path: Path, content: bytes):
    """Write *content* to *path* atomically.

Ignore fsync() issues:
    https://stackoverflow.com/questions/12003805/threadsafe-and-fault-tolerant-file-writes
    """
    with ExitStack() as stack:
        with NamedTemporaryFile(
            dir=str(path.parent), delete=False
        ) as output_file:
            # call cleanup on error in write() or replace()
            def cleanup():
                with suppress(FileNotFoundError):
                    os.remove(output_file.name)

stack.callback(cleanup)

output_file.write(content)

Path(output_file.name).replace(path)
        stack.pop_all()  # success, don't call cleanup