"""
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 .
Copyright © 2019 Cloud Linux Software Inc.
This software is also available under ImunifyAV commercial license,
see
"""
import asyncio
import base64
import errno
import hashlib
import json
import logging
import os
import tempfile
import time
import uuid
from concurrent.futures.process import BrokenProcessPool
from contextlib import asynccontextmanager, suppress
from functools import partial
from uuid import uuid4
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import (
AsyncIterator,
Awaitable,
Callable,
Dict,
List,
Optional,
)
try:
import nats
_has_nats = True
except ImportError: # pragma: no cover - production runtime always has it
nats = None
_has_nats = False
from imav.malwarelib.cleanup import storage_protocol as protocol
from imav.malwarelib.cleanup.types import CleanupRevertInitiator
from imav.malwarelib.config import (
MalwareHitStatus,
MalwareScanResourceType,
)
from imav.malwarelib.model import MalwareHit
from imav.malwarelib.scan.crontab import is_crontab
from imav.malwarelib.tenant_path import split_prefixed, strip_prefix
from imav.malwarelib.scan.mds import restore as mds_restore
from imav.malwarelib.subsys.malware import MalwareAction
from imav.malwarelib.utils import (
hash_path,
publish_with_enospc_reclaim,
verify_failed_restore,
)
from defence360agent.contracts.config import Core
from defence360agent.contracts.config import Malware as Config
from defence360agent.contracts.config import UserType
from defence360agent.internals.the_sink import TheSink
from defence360agent.utils import antivirus_mode, is_cluster, safe_fileops
from defence360agent.subsys.nats_discovery import (
read_nats_addr,
read_nats_token,
)
logger = logging.getLogger(__name__)
#: Hard cap on a file pulled from a syncer into the agent's TMPDIR. The
#: loop trusts the syncer to eventually return eof; a malicious or forged
#: response stream could otherwise fill the agent node's disk.
MAX_STORAGE_FETCH_SIZE = 512 * 1024 * 1024 # 512 MiB
#: Hard cap on a single response file read from the shared volume. The whole
#: file is json.load-ed into memory before use, so an oversized (possibly
#: forged) response is a memory-exhaustion vector. It must stay above the
#: largest legitimate response: a 100 MiB read chunk base64-encodes to
#: ~133 MiB, so 256 MiB leaves headroom for that plus JSON framing.
MAX_RESPONSE_FILE_SIZE = 256 * 1024 * 1024 # 256 MiB
NATS_STORAGE_REQUEST_PREFIX = "imunify360.v1.storage.request."
def _read_first_line(path: str) -> str:
try:
with open(path) as f:
return f.read().strip()
except OSError:
return ""
async def _notify_storage_request_via_nats(
operations: list, request_id: str
) -> None:
"""Publish a NATS signal so the syncer's per-app subscriber wakes
immediately instead of waiting for its fsnotify backup. The request
file on the shared volume is the source of truth — this signal is
a fast-path optimization, so any error here is logged and ignored.
Operations may target multiple application_ids (rare, but legal in a
batch). We publish one message per distinct app_id so each
interested syncer gets the signal exactly once.
"""
if not _has_nats:
return
app_ids: set[str] = set()
for op in operations:
app_id = op.get("application_id") if isinstance(op, dict) else None
if app_id:
app_ids.add(str(app_id))
if not app_ids:
return
addr = read_nats_addr()
try:
token = read_nats_token()
except OSError:
# the broker may be running without auth; a missing file is not fatal
token = ""
connect_kwargs = {
"servers": [f"nats://{addr}"],
"connect_timeout": 2,
"max_reconnect_attempts": 0,
}
if token:
connect_kwargs["token"] = token
try:
nc = await nats.connect(**connect_kwargs)
except Exception as exc:
logger.debug("nats connect for storage signal failed: %s", exc)
return
payload = json.dumps({"request_id": request_id}).encode()
try:
for app_id in app_ids:
try:
await nc.publish(NATS_STORAGE_REQUEST_PREFIX + app_id, payload)
except Exception as exc:
logger.debug(
"nats publish for %s/%s failed: %s",
app_id,
request_id,
exc,
)
try:
await nc.flush(timeout=2)
except Exception:
pass
finally:
try:
await nc.drain()
except Exception:
pass
try:
await nc.close()
except Exception:
pass
class CleanupStoragePublishError(OSError):
"""Raised when a copied original could not be renamed into place.
A plain OSError with ENOENT would be a FileNotFoundError, which the
caller reads as "the malicious file is gone" and deletes the hit.
"""
@dataclass
class RestoreReport:
file: str
scan_id: str
owner: str
initiator: str = UserType.ROOT
cleaned_at: float = -1
reverted_at: float = -1
hash_before_revert: str = ""
hash_after_revert: str = ""
mtime_before_revert: float = -1
mtime_after_revert: float = -1
size_before_revert: float = -1
size_after_revert: float = -1
partial: bool = False
to_dict = asdict
def get_quarantined_original_k8s(file_path: str) -> list[str]:
"""The quarantined original of ``file_path``, if one is stored.
A cleaned file's infected bytes are no longer in the tenant container,
so they cannot be pulled over the storage protocol. The slot is named
from the stored path, which carries the tenant, so no hit row and no
separate app id are needed — and a file whose hit has already been
removed still resolves.
"""
stored = CleanupStorage.path / CleanupStorage.storage_name(str(file_path))
if not stored.exists():
return []
logger.info("File %s is quarantined at %s", file_path, stored)
return [str(stored)]
async def _fetch_file_to_temp_k8s(
app_id: str, file_path: str
) -> Optional[str]:
"""
Pull file_path via the storage protocol into a local tempfile.
Returns the tempfile path, or None if the file could not be fetched.
Caller is responsible for deleting the returned path.
"""
tmp = tempfile.NamedTemporaryFile(
mode="wb", dir=Core.TMPDIR, delete=False, suffix=".mrs"
)
offset = 0
try:
with tmp as f:
while True:
chunk = await CleanupStorageK8s.read_file(
app_id, file_path, offset=offset
)
data = base64.b64decode(chunk.get("data", ""))
if data:
f.write(data)
offset += len(data)
if offset > MAX_STORAGE_FETCH_SIZE:
raise OSError(
f"Storage-protocol stream for {file_path} "
f"(app_id={app_id}) exceeded "
f"{MAX_STORAGE_FETCH_SIZE} bytes; aborting"
)
elif not chunk.get("eof", True):
# Empty data with eof=False would spin on the same
# offset indefinitely. Treat it as a protocol stall.
logger.warning(
"Storage-protocol returned empty non-EOF chunk "
"for %s (app_id=%s, offset=%d); treating as EOF",
file_path,
app_id,
offset,
)
break
if chunk.get("eof", True):
break
return tmp.name
except FileNotFoundError:
with suppress(OSError):
os.unlink(tmp.name)
return None
except Exception as e:
logger.warning(
"Storage-protocol read failed for %s (app_id=%s): %s",
file_path,
app_id,
e,
)
with suppress(OSError):
os.unlink(tmp.name)
return None
@asynccontextmanager
async def resolve_file_for_upload_k8s(
file_path: str,
app_id: Optional[str] = None,
) -> AsyncIterator[List[str]]:
"""
Yield local paths suitable for MRS upload in K8s.
The tenant is the path's own ``/`` prefix: the quarantined
original if the file was already cleaned, otherwise the live bytes pulled
over the storage protocol. A bare shared path names no tenant and nothing
else on the node can disambiguate it, so it resolves to nothing.
Any temp file materialized by the pull is deleted when the caller's
`async with` block exits, whether the upload succeeded or not.
"""
# Callers may pass a MalwareHitPath wrapper or a peewee FilenameField
# value; coerce once here so the rest of the function can assume str.
file_path = str(file_path)
candidate, _ = split_prefixed(file_path)
# The slot name is derived from the path, so an unregistered prefix
# misses rather than reaching a neighbour's slot: a decommissioned
# application can still have its quarantined file uploaded.
upload_paths = get_quarantined_original_k8s(file_path)
temp_path: Optional[str] = None
try:
if not upload_paths:
# The pull needs an application that is actually there to answer.
resolved_app_id = app_id
if not resolved_app_id:
from imav.malwarelib.utils.user_list import registered_apps
if candidate and candidate in await registered_apps():
resolved_app_id = candidate
if not resolved_app_id:
logger.warning(
"Cannot resolve the tenant of %s for MRS", file_path
)
yield []
return
temp_path = await _fetch_file_to_temp_k8s(
resolved_app_id,
# the syncer wire takes the bare in-container path
strip_prefix(file_path, resolved_app_id),
)
if temp_path:
logger.info(
"Pulled %s via storage protocol for MRS upload"
" (app_id=%s)",
file_path,
resolved_app_id,
)
upload_paths = [temp_path]
yield upload_paths
finally:
if temp_path:
with suppress(OSError):
os.unlink(temp_path)
class CleanupStorageFS:
"""
Store files before cleanup and restore them by request
Original filesystem implementation (renamed from CleanupStorage)
"""
path: Path = Path(Config.CLEANUP_STORAGE)
@classmethod
def contains(cls, path) -> bool:
# normalised because the answer decides whether a hit is dropped, and
# a path is not inside the storage just because it starts with it
return Path(os.path.normpath(os.fsdecode(path))).is_relative_to(
cls.path
)
@staticmethod
async def _copy(src: Path, dst: Path, safe_src=False, safe_dst=False):
await safe_fileops.safe_move(
str(src),
str(dst),
src_unlink=False,
dst_overwrite=True,
safe_src=safe_src,
safe_dst=safe_dst,
)
@staticmethod
async def _copy_in_place(
src: Path, dst: Path, safe_src=False, safe_dst=False
):
await safe_fileops.safe_copy_in_place(
str(src),
str(dst),
safe_src=safe_src,
safe_dst=safe_dst,
)
@classmethod
def storage_name(cls, filename: str) -> str:
"""
Get file name for cleanup storage
:return: file name
"""
# In k8s every tenant shares the same web-root path, so a path-only
# key collides across tenants: one app's restore would return
# another app's file. The stored path names its tenant, so the slot
# is scoped from the path itself. The key keeps the tenant and the
# bare path separate rather than hashing the prefixed form, because
# that is how already-quarantined originals were named.
tenant, bare = split_prefixed(str(filename))
if tenant and is_cluster():
return hash_path(f"{tenant}\x00{bare}")
return hash_path(str(filename))
@classmethod
def get_hit_store_path(cls, hit):
return cls.path / cls.storage_name(hit.orig_file)
@classmethod
async def store(cls, hit):
src = hit.orig_file_path
dst = cls.get_hit_store_path(hit)
safe_src = is_crontab(src)
await publish_with_enospc_reclaim(
partial(cls._publish, src, dst, safe_src), cls.path, src, dst
)
@classmethod
async def _publish(cls, src: Path, dst: Path, safe_src: bool):
# the copy truncates its destination before reading the source
tmp = dst.with_name("{}.{}.tmp".format(dst.name, uuid4().hex))
try:
await cls._copy(src, tmp, safe_src=safe_src, safe_dst=True)
try:
os.replace(tmp, dst)
except OSError as e:
raise CleanupStoragePublishError(
e.errno, "failed to publish {}: {}".format(dst, e)
) from e
finally:
with suppress(OSError):
os.unlink(tmp)
@classmethod
async def store_all(cls, hits):
if not cls.path.exists():
cls.path.mkdir(0o700)
succeeded, not_exist = set(), set()
failed: Dict[MalwareHit, Exception] = {}
for hit in hits:
try:
await cls.store(hit)
succeeded.add(hit)
except FileNotFoundError:
not_exist.add(hit)
except (
OSError,
BrokenProcessPool,
safe_fileops.UnsafeFileOperation,
) as e:
logger.warning(
"Failed to store file before cleanup: %r (%r) -- %s",
str(hit),
hit.orig_file,
e,
)
failed[hit] = e
await MalwareAction.cleanup_failed_store(
path=hit.orig_file,
file_owner=hit.owner,
file_user=hit.user,
signature_id=hit.signature_id,
)
return succeeded, failed, not_exist
@classmethod
async def restore(cls, hit: MalwareHit) -> RestoreReport:
report = RestoreReport(hit.orig_file, hit.scanid_id, hit.user)
src = cls.get_hit_store_path(hit)
dst = hit.orig_file_path
try:
return await cls._restore(report, src, dst)
except (
OSError,
BrokenProcessPool,
safe_fileops.UnsafeFileOperation,
) as exc:
try:
(
report.partial,
report.hash_after_revert,
) = await verify_failed_restore(
dst, src, report.hash_before_revert
)
except (
OSError,
BrokenProcessPool,
safe_fileops.UnsafeFileOperation,
):
logger.warning(
"Could not verify restore outcome for %s",
dst,
exc_info=True,
)
exc.restore_report = report
raise
@classmethod
async def _restore(
cls, report: RestoreReport, src: Path, dst: Path
) -> RestoreReport:
safe_dst = is_crontab(dst)
with suppress(FileNotFoundError):
report.cleaned_at = src.stat().st_mtime
with suppress(FileNotFoundError):
st_before = dst.stat()
report.mtime_before_revert = st_before.st_mtime
report.size_before_revert = st_before.st_size
hash_before = hashlib.sha256(dst.read_bytes()).hexdigest()
report.hash_before_revert = hash_before
if not src.stat().st_size:
# an empty backup is what a failed store leaves behind; restoring
# it would empty the file the user still has and report success
raise OSError(
errno.EINVAL, "empty stored original for {}".format(dst)
)
await cls._copy_in_place(src, dst, safe_src=True, safe_dst=safe_dst)
report.reverted_at = time.time()
with suppress(FileNotFoundError):
st_after = dst.stat()
report.mtime_after_revert = st_after.st_mtime
report.size_after_revert = st_after.st_size
hash_after = hashlib.sha256(dst.read_bytes()).hexdigest()
report.hash_after_revert = hash_after
return report
@classmethod
async def restore_all(
cls, hits: List[MalwareHit], initiator: Optional[str] = None
):
succeeded, failed, not_exist = set(), set(), set()
for hit in hits:
try:
report = await cls.restore(hit)
await MalwareAction.cleanup_restored_original(
path=hit.orig_file,
file_owner=hit.owner,
file_user=hit.user,
signature_id=hit.signature_id,
initiator=initiator or UserType.ROOT,
report=report,
)
succeeded.add(hit)
except FileNotFoundError as e:
# ENOENT means either the destination is gone (reap it) or
# the backup is gone (still retryable -- the cleaned file
# remains, so the hit is still meaningful).
# an unanswerable probe must not reap, nor abort the batch
destination_gone = False
with suppress(OSError):
destination_gone = cls.get_hit_store_path(hit).exists()
if destination_gone:
not_exist.add(hit)
else:
await MalwareAction.cleanup_failed_restore(
path=hit.orig_file,
file_owner=hit.owner,
file_user=hit.user,
signature_id=hit.signature_id,
initiator=initiator or UserType.ROOT,
report=getattr(e, "restore_report", None),
)
logger.warning(
"Failed to restore file: %r -- %s", str(hit), e
)
failed.add(hit)
except (
OSError,
BrokenProcessPool,
safe_fileops.UnsafeFileOperation,
) as e:
await MalwareAction.cleanup_failed_restore(
path=hit.orig_file,
file_owner=hit.owner,
file_user=hit.user,
signature_id=hit.signature_id,
initiator=initiator or UserType.ROOT,
report=getattr(e, "restore_report", None),
)
logger.warning("Failed to restore file: %r -- %s", str(hit), e)
failed.add(hit)
return succeeded, failed, not_exist
@classmethod
async def _clear(cls, path: Path, keep: float) -> bool:
try:
st = path.stat()
if st.st_mtime < keep:
path.unlink()
return True
return False
except FileNotFoundError:
return False
except OSError as e:
logger.warning(f"Failed to clear {path.name}: {e}")
return False
@classmethod
async def clear(cls, keep: float) -> int:
"""
Clear storage
:param keep: keep files after specified timestamp
:return:
"""
cls.path.mkdir(0o700, exist_ok=True)
cleared = 0
for path in cls.path.iterdir():
# one entry that vanishes must not abort the sweep
with suppress(FileNotFoundError):
if await cls._clear(path, keep):
cleared += 1
return cleared
class CleanupStorageK8s(CleanupStorageFS):
"""
Message-based cleanup storage for Autonomous platform.
Flow:
1. Python agent writes request to shared volume
2. Resident agent detects via fsnotify
3. Resident agent notifies application syncer
4. Handler executes action, writes response to shared volume
5. Python agent reads response
Application ID is taken from MalwareHit.user field.
All operations are batched for efficiency.
"""
requests_path: Path = Path(protocol.REQUESTS_DIR)
request_timeout: int = protocol.DEFAULT_REQUEST_TIMEOUT
#: Number of retry attempts for transient timeouts: the syncer can be
#: briefly overloaded (concurrent scans, cold cache on the first scan
#: after a pod restart) and miss the initial request_timeout window.
#: Retrying once is enough to cover those cases without turning a real
#: failure into a long stall.
max_retries: int = 1
@classmethod
async def _send_batch_request(
cls, command_type: str, operations: list
) -> dict:
"""
Build, send, and wait for a batch request. Retries on timeout up
to ``max_retries`` extra attempts — each attempt reuses the same
operations but gets a fresh request_id so the syncer cannot return
a stale response.
"""
last_response: dict = {}
for attempt in range(cls.max_retries + 1):
request_id = str(uuid.uuid4())
request = {
"request_id": request_id,
"command_type": command_type,
"operations": operations,
}
cls.requests_path.mkdir(0o700, exist_ok=True)
request_file = cls.requests_path / protocol.get_request_filename(
request_id
)
# Atomic write: the syncer reads via fsnotify and will happily
# json.Unmarshal a half-written file if we write in place. Write
# to a sibling .tmp and rename atomically so the reader either
# sees the full payload or nothing at all.
tmp_file = request_file.with_suffix(request_file.suffix + ".tmp")
with open(tmp_file, "w") as f:
json.dump(request, f, indent=2)
os.replace(tmp_file, request_file)
logger.info(
"Written request %s to %s (attempt %d/%d)",
request_id,
request_file,
attempt + 1,
cls.max_retries + 1,
)
# NATS signal to the per-app syncer subscriber. Best
# effort — the response-file poll below covers any
# publish/connect failure, so we never block on NATS.
await _notify_storage_request_via_nats(operations, request_id)
last_response = await cls._wait_for_response(
request_id, command_type
)
if last_response.get("results"):
return last_response
if attempt < cls.max_retries:
logger.warning(
"Empty results for %s request (attempt %d/%d); retrying",
command_type,
attempt + 1,
cls.max_retries + 1,
)
return last_response
@classmethod
async def _wait_for_response(
cls, request_id: str, command_type: str
) -> dict:
"""
Wait for response file to appear in shared volume and read it.
Since we know the filename in advance, we can poll for it directly.
Syncer application will write the response after processing.
Uses exponential backoff starting at 1s, capped at 30s with 1.5x multiplier.
"""
request_file = cls.requests_path / protocol.get_request_filename(
request_id
)
response_file = cls.requests_path / protocol.get_response_filename(
request_id
)
start_time = time.time()
# Start sub-second so the happy path (syncer replies in <100ms) isn't
# floored at 1s per request. Cap at 2s — a stat() on a shared volume is
# cheap, and extending the interval past 2s mainly adds tail latency on
# top of transient syncer stalls without meaningfully reducing load.
current_interval = 0.05
max_interval = 2.0
multiplier = 1.5
while time.time() - start_time < cls.request_timeout:
if response_file.exists():
try:
if response_file.stat().st_size > MAX_RESPONSE_FILE_SIZE:
raise OSError(
f"response file for {request_id} exceeds "
f"{MAX_RESPONSE_FILE_SIZE} bytes; rejecting"
)
with open(response_file, "r") as f:
response_data = json.load(f)
cls._cleanup_files(request_file, response_file)
logger.info(f"Received response for {request_id}")
return response_data
except (json.JSONDecodeError, OSError) as e:
logger.warning(f"Error reading response file: {e}")
await asyncio.sleep(0.1)
continue
await asyncio.sleep(current_interval)
current_interval = min(current_interval * multiplier, max_interval)
# Force a parent-directory readdir to bypass the NFS
# negative-entry attribute cache. On RWX PVCs (GKE
# Filestore, AWS EFS) a stat() against a path the local
# client has never seen sits behind a ~30 s cache; opendir
# → readdir on the parent invalidates that entry as soon
# as the server returns the new file.
try:
with os.scandir(cls.requests_path) as it:
for _ in it:
pass
except OSError:
pass
# Timeout. Unlink BOTH request and any late response file so a slow
# attempt-1 response doesn't pile up after the client has moved on.
cls._cleanup_files(request_file, response_file)
logger.error(
"Timeout waiting for response request_id=%s command=%s"
" waited=%.1fs",
request_id,
command_type,
time.time() - start_time,
)
return {
"request_id": request_id,
"command_type": command_type,
"results": [],
}
@staticmethod
def _cleanup_files(*files: Path) -> None:
for file in files:
with suppress(OSError):
file.unlink()
@classmethod
def _group_hits_by_app(cls, hits: List[MalwareHit]) -> dict:
"""Group hits by application_id (hit.user)."""
grouped = {}
for hit in hits:
grouped.setdefault(hit.user, []).append(hit)
return grouped
@classmethod
def _build_store_operation(
cls, hit: MalwareHit
) -> protocol.StoreOperation:
src = hit.orig_file_path
dst = cls.get_hit_store_path(hit)
return {
"application_id": hit.user,
"source_path": str(src),
"storage_name": dst.name,
"safe_source": is_crontab(src),
}
@classmethod
def _build_restore_operation(
cls, hit: MalwareHit
) -> protocol.RestoreOperation:
src = cls.get_hit_store_path(hit)
dst = hit.orig_file_path
return {
"application_id": hit.user,
"storage_name": src.name,
"destination_path": str(dst),
"safe_destination": is_crontab(dst),
}
@classmethod
def _populate_report_from_result(
cls, report: RestoreReport, result: dict
) -> None:
"""Populate RestoreReport with before/after stats from result."""
data = result.get("data", {})
if not data:
return
for prefix, attr_prefix in [("before", "before"), ("after", "after")]:
stats = data.get(prefix, {})
if stats:
setattr(
report,
f"mtime_{attr_prefix}_revert",
stats.get("mtime", -1),
)
setattr(
report,
f"size_{attr_prefix}_revert",
stats.get("size", -1),
)
setattr(
report, f"hash_{attr_prefix}_revert", stats.get("hash", "")
)
@classmethod
def _is_not_found_error(cls, error: str) -> bool:
error_lower = error.lower()
return "not found" in error_lower or "does not exist" in error_lower
@classmethod
async def _process_batch(
cls,
app_id: str,
hits: List[MalwareHit],
command_type: str,
build_operation: Callable[[MalwareHit], dict],
on_success: Callable[[MalwareHit, dict], Awaitable[None]],
on_failure: Callable[[MalwareHit, str], Awaitable[None]],
on_not_found: Optional[Callable[[MalwareHit], None]] = None,
) -> tuple[set, set, set]:
"""
Generic batch processor for store/restore operations.
Returns (succeeded, failed, not_exist) sets.
"""
succeeded, failed, not_exist = set(), set(), set()
# Build operations and index hits by storage_name
operations = []
hit_index = {}
for hit in hits:
op = build_operation(hit)
operations.append(op)
hit_index[op["storage_name"]] = hit
response = await cls._send_batch_request(command_type, operations)
results = response.get("results", [])
# Handle timeout (no results)
if not results:
logger.error(
f"No results received for {command_type} batch (app {app_id})"
)
for hit in hits:
failed.add(hit)
await on_failure(hit, "No response received")
return succeeded, failed, not_exist
# Process each result
for result in results:
storage_name = result.get("storage_name")
hit = hit_index.get(storage_name)
if not hit:
logger.warning(
f"Unknown storage_name in response: {storage_name}"
)
continue
if result.get("success"):
await on_success(hit, result)
succeeded.add(hit)
else:
error = result.get("error", "Unknown error")
if on_not_found and cls._is_not_found_error(error):
not_exist.add(hit)
on_not_found(hit)
else:
failed.add(hit)
await on_failure(hit, error)
# A partial response leaves hits in none of the three buckets, so the
# caller would neither retry them nor take them out of their
# in-progress status — they would sit there until something else
# reaped them. Fail them explicitly instead.
for hit in hits:
if (
hit not in succeeded
and hit not in failed
and hit not in not_exist
):
logger.warning(
"No %s result for %r (app %s)",
command_type,
str(hit),
app_id,
)
failed.add(hit)
await on_failure(hit, "No result in response")
return succeeded, failed, not_exist
@classmethod
async def store(cls, hit):
"""
Send STORE request to application side.
Flow:
1. Go resident will get the event and send notification
to application side
2. Syncer read request from NFS
3. Syncer copies file from appFS to shared quarantine volume
(the original stays in place — the cleaner mutates it later)
4. Syncer writes response to shared volume
The retry in _send_batch_request is safe for store because the
syncer's copy is non-destructive and deduped by storage_name, so
a retried request whose first attempt completed returns the
cached success instead of re-copying.
"""
operation = cls._build_store_operation(hit)
response = await cls._send_batch_request("store", [operation])
results = response.get("results", [])
if not results or not results[0].get("success"):
error = (
results[0].get("error", "Unknown error")
if results
else "No results"
)
raise OSError(
f"Store failed: {error} "
f"(storage_name={operation.get('storage_name')!r}, "
f"app_id={hit.user!r}, hit_id={hit.id})"
)
@classmethod
async def store_all(cls, hits):
"""Batch store multiple hits, grouped by application."""
cls.requests_path.mkdir(0o700, exist_ok=True)
all_succeeded, all_not_exist = set(), set()
# dict, matching CleanupStorageFS.store_all: callers read .items()
all_failed: Dict[MalwareHit, Exception] = {}
async def on_success(hit, result):
pass # Nothing extra needed for store success
async def on_failure(hit, error):
logger.warning(
"Failed to store file before cleanup: %r -- %s",
str(hit),
error,
)
all_failed[hit] = OSError(str(error))
await MalwareAction.cleanup_failed_store(
path=hit.orig_file,
file_owner=hit.owner,
file_user=hit.user,
signature_id=hit.signature_id,
)
for app_id, app_hits in cls._group_hits_by_app(hits).items():
succeeded, failed, not_exist = await cls._process_batch(
app_id=app_id,
hits=app_hits,
command_type="store",
build_operation=cls._build_store_operation,
on_success=on_success,
on_failure=on_failure,
on_not_found=lambda h: None,
)
all_succeeded |= succeeded
all_not_exist |= not_exist
return all_succeeded, all_failed, all_not_exist
@classmethod
async def restore(cls, hit: MalwareHit) -> RestoreReport:
"""Send RESTORE request to application side.
Flow:
1. Go resident will get the event and send notification
to application side
2. Syncer read request from NFS
3. Syncer copies file from shared quarantine volume to appFS
4. Syncer deletes file copy from NFS
5. Syncer writes response to shared volume with stat before and after
action
"""
report = RestoreReport(hit.orig_file, hit.scanid_id, hit.user)
# Get cleaned_at from quarantine file (direct FS access)
src = cls.get_hit_store_path(hit)
with suppress(FileNotFoundError):
report.cleaned_at = src.stat().st_mtime
operation = cls._build_restore_operation(hit)
response = await cls._send_batch_request("restore", [operation])
report.reverted_at = time.time()
results = response.get("results", [])
if not results:
raise OSError(
"No results in restore response "
f"(storage_name={operation.get('storage_name')!r}, "
f"app_id={hit.user!r}, hit_id={hit.id})"
)
result = results[0]
if not result.get("success"):
raise OSError(
f"Restore failed: {result.get('error', 'Unknown')} "
f"(storage_name={operation.get('storage_name')!r}, "
f"app_id={hit.user!r}, hit_id={hit.id})"
)
cls._populate_report_from_result(report, result)
return report
@classmethod
async def restore_all(
cls, hits: List[MalwareHit], initiator: Optional[str] = None
):
"""Batch restore multiple hits, grouped by application."""
all_succeeded, all_failed, all_not_exist = set(), set(), set()
reverted_at = time.time()
async def on_success(hit, result):
report = RestoreReport(hit.orig_file, hit.scanid_id, hit.user)
report.reverted_at = reverted_at
src = cls.get_hit_store_path(hit)
with suppress(FileNotFoundError):
report.cleaned_at = src.stat().st_mtime
cls._populate_report_from_result(report, result)
await MalwareAction.cleanup_restored_original(
path=hit.orig_file,
file_owner=hit.owner,
file_user=hit.user,
signature_id=hit.signature_id,
initiator=initiator or UserType.ROOT,
report=report,
)
async def on_failure(hit, error):
logger.warning("Failed to restore file: %r -- %s", str(hit), error)
await MalwareAction.cleanup_failed_restore(
path=hit.orig_file,
file_owner=hit.owner,
file_user=hit.user,
signature_id=hit.signature_id,
initiator=initiator or UserType.ROOT,
)
for app_id, app_hits in cls._group_hits_by_app(hits).items():
# Unlike store, a "not found" syncer error for restore cannot be
# safely attributed to the destination alone -- the syncer's
# error text does not say which side is missing, and a missing
# backup must not be reaped (see CleanupStorageFS.restore_all's
# src/dst state check). Leaving on_not_found unset here means
# every such error stays in `failed`, matching the FS path's
# behavior for the backup-gone case; classifying a genuinely
# gone destination on a cluster is left to its own fix.
succeeded, failed, not_exist = await cls._process_batch(
app_id=app_id,
hits=app_hits,
command_type="restore",
build_operation=cls._build_restore_operation,
on_success=on_success,
on_failure=on_failure,
)
all_succeeded |= succeeded
all_failed |= failed
all_not_exist |= not_exist
return all_succeeded, all_failed, all_not_exist
@classmethod
async def read_file(
cls, app_id: str, path: str, offset: int = 0, limit: int = 65536
) -> dict:
"""
Read a chunk of a file from application storage via syncer.
Returns dict with keys: data (base64 str), size (int), eof (bool).
Raises FileNotFoundError or OSError on failure.
"""
operation = {
"application_id": app_id,
"source_path": path,
"offset": offset,
"limit": limit,
}
response = await cls._send_batch_request("read", [operation])
results = response.get("results", [])
if not results:
raise OSError(
f"No response for read request on {path} (app_id={app_id!r})"
)
result = results[0]
if not result.get("success"):
error = result.get("error", "Unknown")
if cls._is_not_found_error(error):
raise FileNotFoundError(error)
raise OSError(
f"Read failed: {error} (path={path!r}, app_id={app_id!r})"
)
return {
"data": result.get("data", ""),
"size": result.get("size", 0),
"eof": result.get("eof", True),
}
@classmethod
async def stat_file(cls, app_id: str, path: str) -> dict:
"""
Stat a file in application storage via the syncer.
Returns dict with keys: size, mtime, mode, is_file, is_dir.
Raises FileNotFoundError for a missing path, OSError on failure.
"""
operation = {"application_id": app_id, "source_path": path}
response = await cls._send_batch_request("stat", [operation])
results = response.get("results", [])
if not results:
raise OSError(
f"No response for stat request on {path} (app_id={app_id!r})"
)
result = results[0]
if not result.get("success"):
error = result.get("error", "Unknown")
if cls._is_not_found_error(error):
raise FileNotFoundError(error)
raise OSError(
f"Stat failed: {error} (path={path!r}, app_id={app_id!r})"
)
return {
"size": result.get("size", 0),
"mtime": result.get("mtime", 0.0),
"mode": result.get("mode", 0),
"is_file": result.get("is_file", False),
"is_dir": result.get("is_dir", False),
}
@classmethod
async def list_dir(cls, app_id: str, path: str) -> list:
"""
List immediate children of a directory in application storage via the
syncer (non-recursive). Returns a list of entry dicts with keys
name/is_file/is_dir/size/mtime. Raises FileNotFoundError for a missing
directory (or a non-directory path), OSError on failure.
"""
operation = {"application_id": app_id, "source_path": path}
response = await cls._send_batch_request("list-dir", [operation])
results = response.get("results", [])
if not results:
raise OSError(
f"No response for list-dir request on {path} "
f"(app_id={app_id!r})"
)
result = results[0]
if not result.get("success"):
error = result.get("error", "Unknown")
if cls._is_not_found_error(error):
raise FileNotFoundError(error)
raise OSError(
f"List-dir failed: {error} (path={path!r}, app_id={app_id!r})"
)
return result.get("entries") or []
@classmethod
async def clean_preview(cls, app_id: str, path: str) -> dict:
"""
Request syncer to clean a file and return both versions.
The syncer copies the file, runs procu2.php on the copy,
and returns base64-encoded content of both original and cleaned.
Returns dict with keys: original (base64 str), cleaned (base64 str).
Raises FileNotFoundError or OSError on failure.
"""
operation = {
"application_id": app_id,
"source_path": path,
}
response = await cls._send_batch_request("clean-preview", [operation])
results = response.get("results", [])
if not results:
raise OSError(
f"No response for clean-preview request on {path} "
f"(app_id={app_id!r})"
)
result = results[0]
if not result.get("success"):
error = result.get("error", "Unknown")
if cls._is_not_found_error(error):
raise FileNotFoundError(error)
raise OSError(
f"Clean preview failed: {error} "
f"(path={path!r}, app_id={app_id!r})"
)
return {
"original": result.get("original", ""),
"cleaned": result.get("cleaned", ""),
}
CleanupStorage = CleanupStorageK8s if is_cluster() else CleanupStorageFS
async def restore_hits(
hits: list[MalwareHit],
sink: TheSink,
initiator: CleanupRevertInitiator | None = None,
):
file_hits = [
hit
for hit in hits
if hit.resource_type == MalwareScanResourceType.FILE.value
]
succeeded, failed, not_exist = await CleanupStorage.restore_all(
file_hits, initiator
)
MalwareHit.set_status(succeeded, MalwareHitStatus.FOUND)
if antivirus_mode.disabled:
await mds_restore.restore_hits(hits, sink, initiator)
# FIXME: we cannot include db hits here
return succeeded, failed, not_exist