"""
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/>.
Copyright © 2019 Cloud Linux Software Inc.
This software is also available under ImunifyAV commercial license,
see <https://www.imunify360.com/legal/eula>
"""
from __future__ import annotations
import asyncio
import re
import time
from asyncio import CancelledError
from logging import getLogger
from defence360agent.api import inactivity
from defence360agent.contracts.config import Malware as Config
from defence360agent.contracts.license import LicenseError
from defence360agent.contracts.messages import GeneralMetrics, MessageType
from defence360agent.contracts.plugins import (
MessageSink,
MessageSource,
expect,
)
from defence360agent.utils import (
create_task_and_log_exceptions,
safe_cancel_task,
)
from imav.contracts.messages import MalwareMRSUpload
from imav.malwarelib.utils import malware_response
from imav.malwarelib.utils.persistent_mrs_queue import PersistentMRSQueue
logger = getLogger(__name__)
METRICS_REPORT_INTERVAL = 300 # seconds
class MRSUploader(MessageSink, MessageSource):
ERR_MSG = "Failed to submit a file"
SUSP_PATTERN = re.compile(r"(?:suspicious\..+|[CS]MW-SUS-.+|SMW-HEUR-ELF)")
def __init__(self):
self._queue = PersistentMRSQueue()
self._new_items = asyncio.Event()
self._shutdown = asyncio.Event()
async def create_source(self, loop, sink):
self._sink = sink
self._loop = loop
self._queue.cleanup_exhausted()
self._queue.requeue_stale_processing(older_than_seconds=0)
if self._queue.pending_count():
self._new_items.set()
self._upload_task = create_task_and_log_exceptions(loop, self.upload)
await self._maybe_emit_stale_metrics()
self._metrics_task = create_task_and_log_exceptions(
loop, self._report_metrics
)
async def create_sink(self, loop):
pass
async def shutdown(self):
self._shutdown.set()
self._new_items.set()
await safe_cancel_task(self._upload_task)
await self._emit_metrics()
await safe_cancel_task(self._metrics_task)
async def _maybe_emit_stale_metrics(self):
"""Emit metrics on startup if last report is older than interval."""
last_report = self._queue.get_last_metrics_report_time()
if last_report is None or (
time.time() - last_report > METRICS_REPORT_INTERVAL
):
await self._emit_metrics()
async def _report_metrics(self):
"""Periodically emit queue metrics as GENERAL_METRICS."""
while not self._shutdown.is_set():
try:
await asyncio.wait_for(
self._shutdown.wait(),
timeout=METRICS_REPORT_INTERVAL,
)
except asyncio.TimeoutError:
pass
if not self._shutdown.is_set():
await self._emit_metrics()
async def _emit_metrics(self):
"""Collect and emit all queue metrics."""
try:
counters = self._queue.get_counters()
gauges = self._queue.get_gauges()
metrics = {**counters, **gauges}
if not any(metrics.values()):
return
payload = [
{"name": name, "value": value}
for name, value in metrics.items()
]
await self._sink.process_message(GeneralMetrics(payload))
self._queue.set_last_metrics_report_time(time.time())
except Exception:
logger.exception("Failed to emit MRS queue metrics")
def _separate_hits_by_type(self, results) -> tuple:
malicious = []
suspicious = []
extended_suspicious = []
for file, data in results.items():
is_extended_suspicious = False
is_suspicious = False
is_malicious = False
for hit in data["hits"]:
is_extended_suspicious |= hit.get("extended_suspicious", False)
is_suspicious |= bool(
hit["suspicious"]
and self.SUSP_PATTERN.match(hit["matches"])
)
is_malicious |= not hit["suspicious"]
hit_info = malware_response.HitInfo(file, data["hash"])
if is_extended_suspicious:
extended_suspicious.append(hit_info)
elif is_suspicious:
suspicious.append(hit_info)
elif is_malicious:
malicious.append(hit_info)
return malicious, suspicious, extended_suspicious
@expect(MessageType.MalwareScan)
async def process_scan(self, message):
results = message["results"]
if results is None:
return
if not Config.SEND_FILES:
logger.info("Uploading files to MRS is disabled")
return
(
malicious_hits,
suspicious_hits,
extended_suspicious_hits,
) = self._separate_hits_by_type(results)
if malicious_hits:
await self._sink.process_message(
MalwareMRSUpload(
hits=malicious_hits, upload_reason="malicious"
)
)
if suspicious_hits:
# Move uploading to another task
await self._sink.process_message(
MalwareMRSUpload(
hits=suspicious_hits, upload_reason="suspicious"
)
)
if extended_suspicious_hits: # pragma: no cover
await self._sink.process_message(
MalwareMRSUpload(
hits=extended_suspicious_hits,
upload_reason="extended-suspicious",
)
)
errors = message["summary"].get("errors")
if errors:
error_hits = [
malware_response.HitInfo(hit["file"], hit["hash"])
for hit in errors
]
await self._sink.process_message(
MalwareMRSUpload(hits=error_hits, upload_reason="scan_error")
)
async def upload(self):
while not self._shutdown.is_set():
await self._new_items.wait()
self._new_items.clear()
try:
await self._consume_pending()
except CancelledError:
raise
except Exception:
logger.exception("Unexpected error in MRS upload consumer")
if not self._shutdown.is_set():
# A failure mid-batch leaves rows stuck in PROCESSING;
# return them to PENDING so the next cycle picks them up
# instead of idling on _new_items until a restart.
self._queue.requeue_stale_processing(older_than_seconds=0)
self._new_items.set()
await asyncio.sleep(1)
async def _consume_pending(self):
while not self._shutdown.is_set():
pending_reasons = self._queue.pending_reasons()
if not pending_reasons:
return
for upload_reason in pending_reasons:
if self._shutdown.is_set():
return
while not self._shutdown.is_set():
batch = self._queue.take_batch(
upload_reason,
batch_size=malware_response.DEFAULT_CHUNK_SIZE,
)
if not batch:
break
if await self._process_batch(batch, upload_reason):
return
async def _process_batch(self, batch, upload_reason: str):
unknown_hashes = set()
saw_known_hashes_result = False
async for chunk in malware_response.check_known_hashes(
self._loop,
(item.hash for item in batch),
upload_reason,
chunk_size=malware_response.DEFAULT_CHUNK_SIZE,
):
saw_known_hashes_result = True
unknown_hashes.update(chunk)
if not saw_known_hashes_result:
# Generator yielded nothing — `check_known_hashes` short-circuited
# before issuing any request (IAIDTokenError or similar systemic
# failure). Penalising per-item attempts here would EXHAUST the
# whole queue after ~max_attempts seconds of transient auth/CAS
# downtime; treat it like the LicenseError branch and requeue.
self._queue.requeue(
[item.id for item in batch],
"Failed to check whether hashes are already known to MRS",
)
self._new_items.set()
await asyncio.sleep(1)
return True
known_item_ids = [
item.id for item in batch if item.hash not in unknown_hashes
]
if known_item_ids:
self._queue.mark_done(known_item_ids)
unknown_items = [item for item in batch if item.hash in unknown_hashes]
if not unknown_items:
logger.info("All files are known to MRS. Skipping uploading...")
return False
return await self._upload_unknown_items(unknown_items, upload_reason)
async def _upload_unknown_items(self, items, upload_reason: str):
uploaded_item_ids = []
retry_requested = False
with inactivity.track.task("mrs_upload"):
for index, item in enumerate(items):
if self._shutdown.is_set():
if uploaded_item_ids:
self._queue.mark_done(uploaded_item_ids)
self._queue.requeue(
[pending_item.id for pending_item in items[index:]]
)
return True
try:
await malware_response.upload_file(
item.file_path, upload_reason=upload_reason
)
uploaded_item_ids.append(item.id)
except LicenseError as e:
logger.warning(
"Cannot upload files to MRS with reason %s: %s",
upload_reason,
e,
)
if uploaded_item_ids:
self._queue.mark_done(uploaded_item_ids)
self._queue.requeue(
[pending_item.id for pending_item in items[index:]],
str(e),
)
self._new_items.set()
await asyncio.sleep(5)
return True
except malware_response.FileTooLargeError as e:
logger.warning(
"%s: %s (file: %s)",
self.ERR_MSG,
e,
item.file_path,
)
self._queue.mark_failed([item.id], str(e), max_attempts=1)
except malware_response.UploadFailure as e:
logger.warning(
"%s: %s (file: %s)",
self.ERR_MSG,
e,
item.file_path,
)
self._queue.mark_failed([item.id], str(e))
retry_requested = True
except FileNotFoundError as e:
err = "{}. {}".format(self.ERR_MSG, e.strerror)
logger.warning("%s: %s", err, e.filename)
self._queue.mark_failed([item.id], err, max_attempts=1)
except Exception as e:
logger.exception(
"Unexpected MRS upload error for %s",
item.file_path,
)
self._queue.mark_failed([item.id], str(e))
retry_requested = True
if uploaded_item_ids:
self._queue.mark_done(uploaded_item_ids)
if retry_requested:
self._new_items.set()
await asyncio.sleep(1)
return retry_requested
@expect(MessageType.MalwareMRSUpload)
async def process_hits(self, message: MalwareMRSUpload):
upload_reason = message.get("upload_reason", "suspicious")
inserted = self._queue.put_many(message["hits"], upload_reason)
if inserted == 0:
return
self._new_items.set()
|