/opt/imunify360/venv/lib/python3.11/site-packages/imav/contracts
Edit: /opt/imunify360/venv/lib/python3.11/site-packages/imav/contracts/messages.py (11941B)
"""
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 os
from copy import deepcopy
from typing import List, Optional, Set
from defence360agent.contracts.messages import (
Accumulatable,
Ack,
FilesUpdated,
Lockable,
Message,
Noop,
Received,
estimate_size,
Reportable,
ShortenReprListMixin,
Splittable,
)
from imav.malwarelib.config import QueuedScanState
from imav.malwarelib.scan.mds.report import (
MalwareDatabaseHitInfo,
MalwareDatabaseScanReport,
)
class MalwareScanTask(Message):
"""
Creates task to scan file for malware
"""
DEFAULT_METHOD = "MALWARE_SCAN_TASK"
class MalwareScanComplete(Message):
DEFAULT_METHOD = "MALWARE_SCAN_COMPLETE"
class MalwareCleanComplete(Message):
DEFAULT_METHOD = "MALWARE_CLEAN_COMPLETE"
class MalwareRestoreComplete(Message):
DEFAULT_METHOD = "MALWARE_RESTORE_COMPLETE"
# Upper bound on the number of entries kept in unbounded summary arrays
# (errors/stats). A summary carries no per-file results, so this is the only
# thing that can inflate it past the transport payload limit on huge scans.
MALWARE_SUMMARY_MAX_ARRAY = int(
os.environ.get("IMUNIFY360_MALWARE_SUMMARY_MAX_ARRAY", 1000)
)
class MalwareScanSummary(ShortenReprListMixin, Message, Reportable):
DEFAULT_METHOD = "MALWARE_SCAN_SUMMARY"
def __init__(self, *args, items, **kwargs):
# keep the summary only without results
items = [{"summary": self._trim(item["summary"])} for item in items]
super().__init__(*args, items=items, **kwargs)
@staticmethod
def _trim(summary):
trimmed = None
for field in ("errors", "stats"):
value = summary.get(field)
if isinstance(value, list) and len(value) > (
MALWARE_SUMMARY_MAX_ARRAY
):
if trimmed is None:
trimmed = dict(summary)
trimmed[field] = value[:MALWARE_SUMMARY_MAX_ARRAY]
trimmed[f"{field}_truncated"] = len(value)
return trimmed if trimmed is not None else summary
class MalwareScanResult(ShortenReprListMixin, Message, Reportable, Splittable):
DEFAULT_METHOD = "MALWARE_SCAN_RESULT"
LIST_SIZE = 5
BATCH_SIZE = 1000
BATCH_FIELD = "results"
def __init__(self, *args, items, **kwargs):
items = [
{
# keep summary fields necessary for storing results only
"summary": {
key: value
for key, value in item["summary"].items()
if key in {"scanid", "type", "args"}
},
# keep results
"results": item["results"],
}
for item in items
]
super().__init__(*args, items=items, **kwargs)
class MalwareScan(Accumulatable, Lockable):
"""
Represents results of single scan
"""
DEFAULT_METHOD = "MALWARE_SCAN"
LIST_CLASS = (MalwareScanSummary, MalwareScanResult)
def do_accumulate(self) -> bool:
results = self.get("results")
return results is not None
class MalwareMRSUpload(Message):
"""
Used to isolate a possibly long running uploading to MRS from the
other MalwareScan handlers.
"""
DEFAULT_METHOD = "MALWARE_MRS_UPLOAD"
class MalwareScanQueue(Message):
DEFAULT_METHOD = "MALWARE_SCAN_QUEUE"
class MalwareScanQueuePut(MalwareScanQueue):
def __init__(self, *, paths, scan_args):
super().__init__(paths=paths, scan_args=scan_args)
class MalwareScanQueueRemove(MalwareScanQueue):
def __init__(self, *, scan_ids):
super().__init__(scan_ids=scan_ids)
class MalwareScanQueueRecheck(MalwareScanQueue):
def __init__(self):
super().__init__()
class MalwareScanQueueStopBackground(MalwareScanQueue):
def __init__(self):
super().__init__()
class MalwareScanQueueUpdateStatus(MalwareScanQueue):
def __init__(
self, *, scan_ids: List[str] = None, status: QueuedScanState = None
):
super().__init__(scan_ids=scan_ids, status=status)
class MalwareResponse(Message, Received):
TYPES = BLACK, WHITE, KNOWN, UNKNOWN = "BLACK", "WHITE", "KNOWN", "UNKNOWN"
DEFAULT_METHOD = "MALWARE_RESPONSE"
class MalwareSendFiles(Message, Received):
DEFAULT_METHOD = "MALWARE_SEND_FILES"
class MalwareRescanFiles(Message, Received):
DEFAULT_METHOD = "MALWARE_RESCAN_FILES"
class MalwareCleanupTask(Message):
"""
Creates task to cleanup files
"""
DEFAULT_METHOD = "MALWARE_CLEANUP_TASK"
def __init__(
self,
*,
hits,
standard_only=None,
cause=None,
initiator=None,
scan_id=None,
post_action=None,
):
super().__init__(
hits=hits,
standard_only=standard_only,
cause=cause,
initiator=initiator,
scan_id=scan_id,
post_action=post_action,
)
class MalwareCleanupList(
ShortenReprListMixin, Message, Reportable, Splittable
):
DEFAULT_METHOD = "MALWARE_CLEANUP_LIST"
LIST_SIZE = 5
BATCH_SIZE = 1000
BATCH_FIELD = "hits"
@classmethod
def _unit_size(cls, unit, is_dict, message):
# Each hit carries a paired `result[orig_file]` entry (orig_file,
# scan_path, before/after hashes). Count it so the byte budget covers
# the whole per-hit payload, not just the hit itself.
size = super()._unit_size(unit, is_dict, message)
result = message.get("result") or {}
orig_file = getattr(unit, "orig_file", None)
if isinstance(result, dict) and orig_file in result:
size += estimate_size({orig_file: result[orig_file]})
return size
@classmethod
def _split_items(cls, messages: List[Accumulatable]):
for message in super()._split_items(messages):
result = deepcopy(message.get("result", {}))
if result and isinstance(result, dict):
current_paths = {
hit.orig_file for hit in message.get("hits", [])
}
# after splitting messages, leave only the actual result
# to avoid duplicates on the correlation server
extra_paths = result.keys() - current_paths
for path in extra_paths:
result.pop(path)
message["result"] = result
yield message
class MalwareCleanup(Accumulatable):
"""
Represents results of single cleanup
"""
DEFAULT_METHOD = "MALWARE_CLEANUP"
LIST_CLASS = MalwareCleanupList
class MalwareIgnorePathUpdated(Message):
"""Signal through a message bus that MalwareIgnorePath has been updated."""
pass
class MalwareDatabaseScan(Message):
"""
Represents results of a single MDS scan
"""
args: Optional[List[str]]
path: str
scan_id: str
type: Optional[str]
error: Optional[str]
started: int
completed: int
total_resources: int
total_malicious: int
hits: Set[MalwareDatabaseHitInfo]
initiator: Optional[str]
def __init__(
self,
*,
args: Optional[List[str]],
path: Optional[str],
scan_id: str,
type: Optional[str],
started: int = 0,
completed: int = 0,
total_resources: int = 0,
total_malicious: int = 0,
hits: Optional[Set[MalwareDatabaseHitInfo]] = None,
error: Optional[str] = None,
initiator: Optional[str] = None,
):
super().__init__(
args=args,
path=path,
scan_id=scan_id,
type=type,
started=started,
completed=completed,
total_resources=total_resources,
total_malicious=total_malicious,
hits=hits,
error=error,
initiator=initiator,
)
def update_with_report(self, report: MalwareDatabaseScanReport):
self.update(
started=report.started,
completed=report.completed,
total_resources=report.total_resources,
total_malicious=report.total_malicious,
hits=report.hits,
)
def update_with_error(self, message):
return self.update(error=message)
def __setitem__(self, key, value):
raise NotImplementedError
def __delitem__(self, key):
raise NotImplementedError
class MalwareDatabaseRestore(Message):
"""
Represents results of a single MDS restore
"""
class MalwareDatabaseRestoreTask(Message):
"""
Represents a single MDS restore task
"""
path: str
app_name: str
signature_id: str | None
def __init__(self, path: str, app_name: str, signature_id: str | None):
super().__init__(
path=path, app_name=app_name, signature_id=signature_id
)
class MalwareDatabaseCleanup(Message):
"""
Represents results of a single mds cleanup
"""
succeeded: Set[MalwareDatabaseHitInfo]
failed: Set[MalwareDatabaseHitInfo]
def __init__(
self,
succeeded: Set[MalwareDatabaseHitInfo],
failed: Set[MalwareDatabaseHitInfo],
scan_id: Optional[str] = None,
):
super().__init__(
succeeded=succeeded,
failed=failed,
scan_id=scan_id,
)
class MalwareDatabaseCleanupFailed(Message):
"""
Signifies an MDS cleanup failure
"""
error: str
def __init__(self, *, error: str):
super().__init__(error=error)
class MalwareDatabaseRestoreFailed(Message):
"""
Signifies an MDS restore failure
"""
error: str
def __init__(self, *, error: str):
super().__init__(error=error)
class RefreshImunifyPatchSubscription(Message):
"""Refresh Imunify Patch user subscription info."""
class MalwareCleanupRevertList(ShortenReprListMixin, Message, Reportable):
DEFAULT_METHOD = "MALWARE_CLEANUP_REVERT_LIST"
class MalwareCleanupRevert(Accumulatable):
LIST_CLASS = MalwareCleanupRevertList
class CheckDetachedScans(Message):
DEFAULT_METHOD = "MALWARE_CHECK_DETACHED_SCANS"
class VulnerabilityPatchTask(Message):
DEFAULT_METHOD = "VULNERABILITY_PATCH_TASK"
class VulnerabilityPatchList(ShortenReprListMixin, Message, Reportable):
DEFAULT_METHOD = "MALWARE_CLEANUP_LIST"
class VulnerabilityPatch(Accumulatable):
DEFAULT_METHOD = "MALWARE_CLEANUP"
LIST_CLASS = VulnerabilityPatchList
class VulnerabilityPatchRevertList(ShortenReprListMixin, Message, Reportable):
DEFAULT_METHOD = "MALWARE_CLEANUP_REVERT_LIST"
class VulnerabilityPatchRevert(Accumulatable):
LIST_CLASS = VulnerabilityPatchRevertList
class VulnerabilityPatchFailed(Message, Reportable):
DEFAULT_METHOD = "CLEANUP_FAILED"
MSGS_WITHOUT_IP = [
msg.DEFAULT_METHOD
for msg in (
Ack,
CheckDetachedScans,
MalwareScan,
MalwareScanSummary,
MalwareScanResult,
MalwareScanComplete,
MalwareCleanComplete,
MalwareRestoreComplete,
MalwareSendFiles,
Noop,
FilesUpdated,
)
]