22

✏️ Editing: storage_protocol.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>

Message Protocol for Filesystem-Based Storage Requests

This module defines the request/response format for file operations
that are passed through the shared volume.

Flow:
1. Python agent writes <request_id>.request.json to shared volume
2. Python agent publishes a best-effort NATS doorbell
   (imunify360.v1.storage.request.<app_id>)
3. The syncer's storage subscriber picks up the request
4. Handler executes action, writes <request_id>.response.json to shared volume
5. Python agent reads response (filename is predictable)

All operations include application_id (from MalwareHit.user field)

Batch Operations:
- Requests and responses are always lists of operations
- Single operations are wrapped in a one-element list
- This allows efficient batch processing of multiple files
"""
import os
from typing import Literal, TypedDict, Optional, List

class BaseOperation(TypedDict):
    """Base structure for a single operation within a batch"""

application_id: str

class StoreOperation(BaseOperation):
    """
    Single store operation in a batch.

The application side handler should:
    1. Read file from source_path in application storage
    2. Write it to shared quarantine volume with storage_name
    3. Delete original file from application storage
    """

source_path: str  # Path in application (e.g., /var/www/html/malware.php)
    storage_name: str  # Hashed name for quarantine (e.g., a1b2c3d4...)
    safe_source: bool  # Whether source needs safe handling (is crontab)

class RestoreOperation(BaseOperation):
    """
    Single restore operation in a batch.

The application side handler should:
    1. Collect "before" stats of destination file (if exists)
       - mtime, size, hash
    2. Read file from shared quarantine volume (storage_name)
    3. Write it to application storage at destination_path
    4. Collect "after" stats of destination file
       - mtime, size, hash
    5. Return both "before" and "after" stats in response
    """

storage_name: str  # Name in quarantine storage (on shared volume)
    destination_path: str  # Path in application to restore to
    safe_destination: bool  # Whether destination needs safe handling

class ReadOperation(BaseOperation):
    """
    Single read operation in a batch.

The application side handler should:
    1. Read file at source_path in application storage
    2. Read at most `limit` bytes starting from `offset`
    3. Return base64-encoded chunk and file size in response
    """

source_path: str  # Path in application (e.g., /var/www/html/file.php)
    offset: int  # Byte offset to start reading from
    limit: int  # Maximum number of bytes to read

class CleanPreviewOperation(BaseOperation):
    """
    Single clean-preview operation.

The application side handler should:
    1. Read infected file from source_path
    2. Copy to temp, run procu2.php on the copy
    3. Return base64-encoded content of both original and cleaned versions
    """

source_path: str  # Path in application (infected file)

class StatOperation(BaseOperation):
    """
    Single stat operation in a batch.

The application side handler should:
    1. Stat source_path in application storage (following symlinks)
    2. Return type flags, size, mtime and POSIX mode in the response
    3. Report a non-existent path as a not-found failure, so the client
       can map it to FileNotFoundError
    """

source_path: str  # Path in application to stat

class ListDirOperation(BaseOperation):
    """
    Single directory-listing operation in a batch.

The application side handler should:
    1. Read the immediate children of source_path (non-recursive)
    2. Return each entry's base name, type flags, size and mtime
    3. Cap the number of entries (truncated signals the cap was hit)
    """

source_path: str  # Path in application of the directory to list

class BatchRequest(TypedDict):
    """
    Batch request containing multiple operations.

Even single operations are wrapped in a one-element list.
    """

request_id: str
    command_type: Literal[
        "store", "restore", "read", "clean-preview", "stat", "list-dir"
    ]
    operations: (
        List[StoreOperation]
        | List[RestoreOperation]
        | List[ReadOperation]
        | List[CleanPreviewOperation]
        | List[StatOperation]
        | List[ListDirOperation]
    )

class FileStats(TypedDict):
    """File statistics for before/after comparison"""

mtime: float  # Modification time
    size: int  # File size in bytes
    hash: str  # SHA256 hash of file contents

class StoreOperationResult(TypedDict):
    """Result of a single store operation"""

storage_name: str  # Which file this result is for
    success: bool
    error: Optional[str]

class RestoreOperationData(TypedDict):
    """Data for a successful restore operation"""

before: Optional[
        FileStats
    ]  # Stats before restore (None if file didn't exist)
    after: FileStats  # Stats after restore

class RestoreOperationResult(TypedDict):
    """Result of a single restore operation"""

storage_name: str  # Which file this result is for
    success: bool
    error: Optional[str]
    data: Optional[RestoreOperationData]

class ReadOperationResult(TypedDict):
    """Result of a single read operation"""

source_path: str  # Which file this result is for
    success: bool
    error: Optional[str]
    data: Optional[str]  # Base64-encoded file chunk (None on failure)
    size: Optional[int]  # Total file size in bytes (None on failure)
    eof: bool  # True if chunk reached end of file

class CleanPreviewResult(TypedDict):
    """Result of a single clean-preview operation"""

source_path: str  # Which file this result is for
    success: bool
    error: Optional[str]
    original: Optional[str]  # Base64-encoded original (infected) content
    cleaned: Optional[
        str
    ]  # Base64-encoded cleaned content (empty string if removed)

class StatOperationResult(TypedDict):
    """
    Result of a single stat operation.

mtime is Unix seconds with sub-second precision, matching
    os.stat_result.st_mtime; mode is the POSIX st_mode. On a not-found
    path success is False and error contains "not found" so the client
    maps it to FileNotFoundError.
    """

source_path: str  # Which file this result is for
    success: bool
    error: Optional[str]
    size: int  # File size in bytes
    mtime: float  # Modification time (Unix seconds, fractional)
    mode: int  # POSIX st_mode (type bits | permission bits)
    is_file: bool
    is_dir: bool

class ListDirEntry(TypedDict):
    """One child of a listed directory; name is the base name only."""

name: str
    is_file: bool
    is_dir: bool
    size: int
    mtime: float  # Modification time (Unix seconds, fractional)

class ListDirOperationResult(TypedDict):
    """
    Result of a single list-dir operation.

truncated is True when the directory held more entries than the
    syncer's cap — the listing must not be treated as complete.
    """

source_path: str  # Which directory this result is for
    success: bool
    error: Optional[str]
    entries: Optional[List[ListDirEntry]]  # None on failure
    truncated: bool

class BatchResponse(TypedDict):
    """
    Batch response containing results for all operations.

Results are returned in the same order as operations in the request.
    """

request_id: str
    command_type: Literal[
        "store", "restore", "read", "clean-preview", "stat", "list-dir"
    ]
    results: (
        List[StoreOperationResult]
        | List[RestoreOperationResult]
        | List[ReadOperationResult]
        | List[CleanPreviewResult]
        | List[StatOperationResult]
        | List[ListDirOperationResult]
    )

# File naming convention
REQUEST_FILE_PATTERN = "{request_id}.request.json"
RESPONSE_FILE_PATTERN = "{request_id}.response.json"

# File locations (shared volume)
REQUESTS_DIR = os.getenv(
    "I360_CLEANUP_STORAGE_REQUESTS_DIR", "/var/imunify360/.storage_requests"
)

# Timeouts
#: NFS-backed RWX PVCs (GKE Filestore, EFS) cache negative-entry stat
#: results for up to ~30 s by default (acregmax). The syncer writes the
#: response file in milliseconds, but the agent's stat() can return
#: "no such file" until the cache expires. Setting the timeout above the
#: 30 s cache horizon gives at least one post-expiry poll a chance to
#: see the new file. The waiter still short-circuits as soon as it does.
DEFAULT_REQUEST_TIMEOUT = 90  # seconds

def get_request_filename(request_id: str) -> str:
    """Get the filename for a request"""
    return REQUEST_FILE_PATTERN.format(request_id=request_id)

def get_response_filename(request_id: str) -> str:
    """Get the filename for a response"""
    return RESPONSE_FILE_PATTERN.format(request_id=request_id)

MR SKB