22

✏️ Editing: 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
STARTUP_SOCKET_TIMEOUT = 0.1

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(timeout=STARTUP_SOCKET_TIMEOUT):
    """Notify aibolit that there is a file to scan.

Start/restart aibolit if necessary.
    """
    # sending a byte to the aibolit's startup socket starts aibolit if it
    # is not running already

# note: assume that if aibolit is running then it
    # picks up the job itself
    with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as sock:
        sock.settimeout(timeout)
        sock.connect(AIBOLIT_STARTUP_SOCKET)
        sock.send(b"1")

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

MR SKB