From 930966e051b2786e19383d305a788af4dd0c4787 Mon Sep 17 00:00:00 2001 From: su-amaas Date: Mon, 10 Aug 2026 08:46:28 +0000 Subject: [PATCH] update to latest version: v1.4.7 --- CHANGELOG.md | 11 ++++++ README.md | 67 ++++++++++++++++++++++++++++++++++-- VERSION | 2 +- amaas/grpc/__init__.py | 17 ++++++--- amaas/grpc/aio/__init__.py | 14 ++++++-- amaas/grpc/util.py | 10 ++++-- examples/README.md | 2 +- protos/scan.proto | 7 ++-- tests/mock_server.py | 4 ++- tests/test_aio_client_sdk.py | 51 ++++++++++++++++++++++++++- tests/test_client_sdk.py | 50 +++++++++++++++++++++++++++ 11 files changed, 217 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 47629b5..0278003 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,16 @@ # CHANGELOG +## 1.4.7 - 2026-08-10 + +- Support new region ap-southeast-3 (Indonesia) +- Support scan gateway for FSCS and FSVA +- Add an Error Handling section to the README documenting SDK and service error codes and messages + +## 1.4.6 - 2026-06-24 + +- `scan_file` now reports the file path you pass as the scan identifier, instead of only + its basename. Scan results (e.g. `foundMalwares[].fileName`) reflect the full path. + ## 1.4.5 - 2026-03-18 - Support new region af-south-1 diff --git a/README.md b/README.md index 11a525f..2259b4e 100644 --- a/README.md +++ b/README.md @@ -200,7 +200,7 @@ Creates a new instance of the grpc Channel, and provisions essential settings, i | Parameter | Description | | ---------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| region | The region you obtained your api key. Value provided must be one of the Vision One regions, e.g. `us-east-1`, `eu-central-1`, `ap-northeast-1`, `ap-southeast-2`, `ap-southeast-1`, `ap-south-1`, `me-central-1`, `ca-central-1`, `eu-west-2`, `af-south-1`,etc. | +| region | The region you obtained your api key. Value provided must be one of the Vision One regions, e.g. `us-east-1`, `eu-central-1`, `ap-northeast-1`, `ap-southeast-2`, `ap-southeast-1`, `ap-south-1`, `me-central-1`, `ca-central-1`, `eu-west-2`, `af-south-1`, `ap-southeast-3`,etc. | | api_key | Your own Vision One API Key. | | enable_tls | Enable or disable TLS. TLS should always be enabled when connecting to the AMaaS server. For more information, see the 'Ensuring Secure Communication with TLS' section. | | ca_cert | `Optional` CA certificate used to connect to self hosted AMaaS server. | @@ -216,7 +216,7 @@ Creates a new instance of the grpc aio Channel, and provisions essential setting | Parameter | Description | | ---------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| region | The region you obtained your api key. Value provided must be one of the Vision One regions, e.g. `us-east-1`, `eu-central-1`, `ap-northeast-1`, `ap-southeast-2`, `ap-southeast-1`, `ap-south-1`, `me-central-1`, `ca-central-1`, `eu-west-2`, `af-south-1`, etc. | +| region | The region you obtained your api key. Value provided must be one of the Vision One regions, e.g. `us-east-1`, `eu-central-1`, `ap-northeast-1`, `ap-southeast-2`, `ap-southeast-1`, `ap-south-1`, `me-central-1`, `ca-central-1`, `eu-west-2`, `af-south-1`, `ap-southeast-3`, etc. | | api_key | Your own Vision One API Key. | | enable_tls | Enable or disable TLS. TLS should always be enabled when connecting to the AMaaS server. For more information, see the 'Ensuring Secure Communication with TLS' section. | | ca_cert | `Optional` CA certificate used to connect to self hosted AMaaS server. | @@ -286,6 +286,69 @@ Remember to clean up the grpc aio Channel when you are done using it to release | --------- | ------------------------------------------------------------- | | handle | The grpc aio Channel instance created from the init function. | +## Error Handling + +The File Security Python SDK raises `amaas.grpc.exception.AMaasException` for every error condition produced by `amaas.grpc` and `amaas.grpc.aio`. Each exception carries an `error_code` (a member of the `AMaasErrorCode` enum) and a formatted `message`; `str(exception)` renders both as `: `. + +```python +from amaas.grpc.exception import AMaasException + +try: + result = amaas.grpc.scan_file(handle, file_name=filename, tags=tags) +except AMaasException as e: + print(e.error_code, e.message) +except Exception as e: + print(e) +``` + +Some conditions are detected by the SDK itself, independently of any network call: an unsupported region, a missing/unreadable file, invalid tags, or an unexpected message in the scan protocol stream. Everything else comes from the File Security service: the SDK catches `grpc.RpcError` (`grpc.aio.AioRpcError` in the asyncio client) and re-raises it as an `AMaasException`. + +For most service errors the SDK preserves the gRPC status code and message text exactly as sent by the service; for two conditions — authentication failures and rate limiting — the SDK discards the service's own message and substitutes a fixed string of its own. Both the synchronous client and the asyncio client apply the same mapping. + +The **Source** column classifies each error: + +- **SDK-native** — produced entirely by the SDK on the client side, without a network call (unsupported region, missing/unreadable file, invalid tags, unexpected protocol message). +- **SDK-mapped** — triggered by a gRPC response from the service, but the caller-visible message is a fixed string produced by the SDK (authentication failures and rate limiting). +- **Service** — the service's gRPC code and message are relayed to the caller unchanged (via `MSG_ID_GRPC_ERROR`). + +| gRPC status code (as received) | `AMaasErrorCode` / message seen by the caller | Cause | Source | +| --- | --- | --- | --- | +| — | `MSG_ID_ERR_INVALID_REGION`: ` is not a supported region, region value should be one of ` | Region passed to `init_by_region` is not a supported Vision One region | SDK-native | +| — | `MSG_ID_ERR_FILE_NOT_FOUND`: `Failed to open file. No such file or directory .` | File passed to `scan_file` does not exist | SDK-native | +| — | `MSG_ID_ERR_FILE_NO_PERMISSION`: `Failed to open file. Permission denied to open .` | No OS permission to read the file | SDK-native | +| — | `MSG_ID_ERR_INVALID_TAG`: `Invalid tag format: .` | A tag is empty or longer than 63 characters, checked client-side before the scan request is sent | SDK-native | +| — | `MSG_ID_ERR_TAG_NUMBER_EXCEED`: `Too many tags: .` | More than 8 tags supplied, checked client-side before the scan request is sent | SDK-native | +| — | `MSG_ID_ERR_UNKNOWN_CMD` / `MSG_ID_ERR_UNKNOWN_STAGE` / `MSG_ID_ERR_UNEXPECTED_CMD_AND_STAGE`: `Received unknown command from server: ` / `Received unknown stage from server: ` / `Received unexpected command and stage .` | The scan protocol stream produced a command/stage the SDK does not recognize | SDK-native | +| — | `MSG_ID_ERR_UNEXPECTED_ERROR`: `Unexpected error encountered. ` | Any other, non-gRPC exception raised while scanning (also used internally if an unsupported hash algorithm is requested) | SDK-native | +| `Unauthenticated` (16) | `MSG_ID_ERR_KEY_AUTH_FAILED`: `Invalid token or Api Key.` | Service rejected the request as `Unauthenticated` — covers a missing key, an invalid/expired key, and an account without file-scan permission. The SDK substitutes this fixed string for all three cases; the service's actual message is not shown to the caller | SDK-mapped | +| `Internal` (13), details containing `429` ¹ | `MSG_ID_ERR_RATE_LIMIT_EXCEEDED`: `Raised by the SDK library to indicate http 429 too many request error.` | Rate limit exceeded. The SDK detects this by scanning the raw error text for the substring `429` and substitutes this fixed string; the service's actual message is not shown to the caller | SDK-mapped | +| `InvalidArgument` (3) | `MSG_ID_GRPC_ERROR`: `Too many tags. Decrease to eight tags or less.` | Too many tags (only reaches the service if the SDK's own client-side check didn't already catch it) | Service | +| `InvalidArgument` (3) | `MSG_ID_GRPC_ERROR`: `Tag is too long. Decrease length to 63 characters or less.` | A tag longer than 63 characters (only reaches the service if the client-side check didn't already catch it) | Service | +| `InvalidArgument` (3) | `MSG_ID_GRPC_ERROR`: `Tag is empty. Remove the tag or add at least one character.` | An empty tag | Service | +| `InvalidArgument` (3) | `MSG_ID_GRPC_ERROR`: `cloudAccountId contains illegal characters (#, @)` | Illegal characters in a `cloudAccountId` tag | Service | +| `InvalidArgument` (3) | `MSG_ID_GRPC_ERROR`: `Prefix format or length of SHA1 from the SDK is incorrect. Contact Support.` / `SHA1 format from the SDK is incorrect. Contact Support.` | Malformed SHA1 digest sent by the SDK | Service | +| `InvalidArgument` (3) | `MSG_ID_GRPC_ERROR`: `Prefix format or length of SHA256 from the SDK is incorrect. Contact Support.` / `SHA256 format from the SDK is incorrect. Contact Support.` | Malformed SHA256 digest sent by the SDK | Service | +| `NotFound` (5) | `MSG_ID_GRPC_ERROR`: `Customer ID not found. Contact Support.` | Account / customer ID not found | Service | +| `PermissionDenied` (7) | `MSG_ID_GRPC_ERROR`: `SDK feature is not enabled for this account. Contact your administrator to enable the SDK feature.` | The SDK feature is not enabled for the account | Service | +| `ResourceExhausted` (8) | `MSG_ID_GRPC_ERROR`: `This account has performed five scans in the last hour. Purchase and allocate credits to File Security or wait an hour to make five more scans.` | Hourly scan quota exhausted (Essential accounts) | Service | +| `ResourceExhausted` (8) | `MSG_ID_GRPC_ERROR`: `file size is over maximum allowed size ` | Scanned file/buffer exceeds the maximum allowed size | Service | +| `ResourceExhausted` (8) | `MSG_ID_GRPC_ERROR`: `Cannot allocate resource. Try again later. If the issue persists, contact Support.` | Service could not allocate a scan resource | Service | +| `FailedPrecondition` (9) | `MSG_ID_GRPC_ERROR`: `Incorrect stage from the SDK. Contact Support.` | Incorrect protocol stage | Service | +| `Unimplemented` (12) | `MSG_ID_GRPC_ERROR`: `Predictive Machine Learning is not supported. Contact Support.` | PML requested but not supported for the account/region | Service | +| `Internal` (13) | `MSG_ID_GRPC_ERROR`: `Failed to retrieve metadata. Try again later. If the issue persists, contact Support.` | Service could not retrieve request metadata | Service | +| `Internal` (13) | `MSG_ID_GRPC_ERROR`: `Network connection error. Try again later. If the issue persists, contact Support.` | Service-side network / connection error | Service | +| `Internal` (13) | `MSG_ID_GRPC_ERROR`: `Internal error. Try again later. If the issue persists, contact Support.` | Generic internal service error | Service | +| `Internal` (13) | `MSG_ID_GRPC_ERROR`: `Missing preamble information from the SDK. Contact Support.` | Missing preamble information in the scan request | Service | +| `Internal` (13) | `MSG_ID_GRPC_ERROR`: `Unclear scan result: . Contact Support.` | Service could not parse the scan result | Service | +| any other code (preserved) | `MSG_ID_GRPC_ERROR`: `Received gRPC status code: , msg:
.` | Any other error relayed from the service, with the numeric gRPC code and message exactly as the service sent them | Service | + +**Notes** + +1. The service currently signals rate limiting with `Internal` and a details string containing `Http Error Code: 429`; the SDK matches on that substring rather than on a dedicated gRPC code, so this row is reached before the generic `MSG_ID_GRPC_ERROR` case below. +2. Rows marked `MSG_ID_GRPC_ERROR` are relayed from the service with the gRPC status code and message preserved exactly; `str(exception)` renders them as `MSG_ID_GRPC_ERROR: Received gRPC status code: , msg: .`. These messages are owned by the File Security service and may change independently of the SDK. + The Python SDK does not expose Encode/Decode operations, so the service's encode/decode messages cannot occur here. +3. Engine findings such as `ATSE_*` codes are not errors — they are returned inside a successful scan result payload, not raised as an exception. + ## Environment Variables The following environment variables are supported by Python Client SDK and can be used in lieu of values specified as function arguments. diff --git a/VERSION b/VERSION index 03e5161..be05bba 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -1.4.5 \ No newline at end of file +1.4.7 diff --git a/amaas/grpc/__init__.py b/amaas/grpc/__init__.py index bc68ebf..40a916f 100644 --- a/amaas/grpc/__init__.py +++ b/amaas/grpc/__init__.py @@ -23,6 +23,7 @@ logger.propagate = False timeout_in_seconds = int(os.environ.get('TM_AM_SCAN_TIMEOUT_SECS', 300)) +heartbeat_interval_in_seconds = int(os.environ.get('TM_AM_HEARTBEAT_INTERVAL_SECS', 30)) class _Pipeline: @@ -36,8 +37,10 @@ def __init__(self): self._consumer_lock = threading.Lock() self._consumer_lock.acquire() - def get_message(self): - self._consumer_lock.acquire() + def get_message(self, timeout=None): + acquired = self._consumer_lock.acquire(timeout=timeout) + if not acquired: + return None message = self._message self._producer_lock.release() return message @@ -81,7 +84,13 @@ def _generate_messages(pipeline: _Pipeline, data_reader: BinaryIO, bulk: bool, s yield response responses.clear() - message = pipeline.get_message() + while True: + message = pipeline.get_message(timeout=heartbeat_interval_in_seconds) + if message is None: + logger.debug("sending heartbeat to keep connection alive") + yield scan_pb2.C2S(stage=scan_pb2.STAGE_HEARTBEAT) + continue + break if message.stage == scan_pb2.STAGE_INIT: logger.debug("stage INIT") @@ -191,7 +200,7 @@ def scan_file(channel: grpc.Channel, file_name: str, tags: List[str] = None, pml: bool = False, feedback: bool = False, verbose: bool = False, digest: bool = True) -> str: try: f = open(file_name, "rb") - fid = os.path.basename(file_name) + fid = file_name n = os.stat(file_name).st_size except FileNotFoundError as err: logger.debug("File not exist: " + str(err)) diff --git a/amaas/grpc/aio/__init__.py b/amaas/grpc/aio/__init__.py index a673cf3..b491a6d 100644 --- a/amaas/grpc/aio/__init__.py +++ b/amaas/grpc/aio/__init__.py @@ -1,3 +1,4 @@ +import asyncio import io import os from typing import BinaryIO, List @@ -22,6 +23,7 @@ logger.propagate = False timeout_in_seconds = int(os.environ.get('TM_AM_SCAN_TIMEOUT_SECS', 300)) +heartbeat_interval_in_seconds = int(os.environ.get('TM_AM_HEARTBEAT_INTERVAL_SECS', 30)) def init_by_region(region, api_key, enable_tls=True, ca_cert=None): @@ -75,7 +77,15 @@ async def _scan_data(channel: grpc.Channel, data_reader: BinaryIO, size: int, id await call.write(request) while True: - response = await call.read() + while True: + try: + response = await asyncio.wait_for( + call.read(), timeout=heartbeat_interval_in_seconds + ) + break + except asyncio.TimeoutError: + logger.debug("sending heartbeat to keep connection alive") + await call.write(scan_pb2.C2S(stage=scan_pb2.STAGE_HEARTBEAT)) if response.cmd == scan_pb2.CMD_RETR: if response.stage != scan_pb2.STAGE_RUN: @@ -149,7 +159,7 @@ async def scan_file(channel: grpc.Channel, file_name: str, tags: List[str] = Non pml: bool = False, feedback: bool = False, verbose: bool = False, digest: bool = True) -> str: try: f = open(file_name, "rb") - fid = os.path.basename(file_name) + fid = file_name n = os.stat(file_name).st_size except FileNotFoundError as err: logger.debug("File not exist: " + str(err)) diff --git a/amaas/grpc/util.py b/amaas/grpc/util.py index 67f6a91..b8be5df 100644 --- a/amaas/grpc/util.py +++ b/amaas/grpc/util.py @@ -21,6 +21,7 @@ AWS_GB_REGION = "eu-west-2" AWS_AE_REGION = "me-central-1" AWS_ZA_REGION = "af-south-1" +AWS_ID_REGION = "ap-southeast-3" C1_JP_REGION = "jp-1" C1_SG_REGION = "sg-1" C1_AU_REGION = "au-1" @@ -32,13 +33,14 @@ C1_GB_REGION = "gb-1" C1_AE_REGION = "ae-1" C1_ZA_REGION = "za-1" +C1_ID_REGION = "id-1" C1Regions = [C1_AU_REGION, C1_CA_REGION, C1_DE_REGION, C1_GB_REGION, C1_IN_REGION, C1_JP_REGION, C1_SG_REGION, - C1_US_REGION, C1_TREND_REGION, C1_ZA_REGION] -V1Regions = [AWS_AU_REGION, AWS_DE_REGION, AWS_IN_REGION, AWS_JP_REGION, AWS_SG_REGION, AWS_US_REGION, AWS_AE_REGION, AWS_CA_REGION, AWS_GB_REGION, AWS_ZA_REGION] + C1_US_REGION, C1_TREND_REGION, C1_ZA_REGION, C1_ID_REGION] +V1Regions = [AWS_AU_REGION, AWS_DE_REGION, AWS_IN_REGION, AWS_JP_REGION, AWS_SG_REGION, AWS_US_REGION, AWS_AE_REGION, AWS_CA_REGION, AWS_GB_REGION, AWS_ZA_REGION, AWS_ID_REGION] SupportedV1Regions = V1Regions SupportedC1Regions = [C1_AU_REGION, C1_CA_REGION, C1_DE_REGION, C1_GB_REGION, C1_IN_REGION, C1_JP_REGION, C1_SG_REGION, - C1_US_REGION, C1_ZA_REGION] + C1_US_REGION, C1_ZA_REGION, C1_ID_REGION] AllRegions = C1Regions + V1Regions AllValidRegions = SupportedC1Regions + SupportedV1Regions @@ -54,6 +56,7 @@ AWS_CA_REGION: C1_CA_REGION, AWS_GB_REGION: C1_GB_REGION, AWS_ZA_REGION: C1_ZA_REGION, + AWS_ID_REGION: C1_ID_REGION, } @@ -105,6 +108,7 @@ def _init_by_region_util(region, api_key, enable_tls=True, ca_cert=None, is_aio_ C1_CA_REGION: 'antimalware.ca-1.cloudone.trendmicro.com:443', C1_AE_REGION: 'antimalware.ae-1.cloudone.trendmicro.com:443', C1_ZA_REGION: 'antimalware.za-1.cloudone.trendmicro.com:443', + C1_ID_REGION: 'antimalware.id-1.cloudone.trendmicro.com:443', } # make sure it is valid V1 or C1 region diff --git a/examples/README.md b/examples/README.md index 3ea18d5..6423b1a 100644 --- a/examples/README.md +++ b/examples/README.md @@ -56,7 +56,7 @@ If you plan on using a Trend Vision One region, be sure to pass in region parame | ------------------------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -------------------- | | -f FILENAME, --filename FILENAME | File to be scanned | No | | -a ADDR, --addr ADDR | Trend Vision One File Security server | Yes, either -r or -a | - | -r REGION, --region REGION | The region you obtained your API key. Value provided must be one of the Vision One regions, e.g. `us-east-1`, `eu-central-1`, `ap-northeast-1`, `ap-southeast-2`, `ap-southeast-1`, `ap-south-1`, `me-central-1`, `ca-central-1`, `eu-west-2`, `af-south-1` | Yes, either -r or -a | + | -r REGION, --region REGION | The region you obtained your API key. Value provided must be one of the Vision One regions, e.g. `us-east-1`, `eu-central-1`, `ap-northeast-1`, `ap-southeast-2`, `ap-southeast-1`, `ap-south-1`, `me-central-1`, `ca-central-1`, `eu-west-2`, `af-south-1`, `ap-southeast-3` | Yes, either -r or -a | | --api_key API_KEY | Vision One API Key | Yes | | --tls, --no-tls | Enable or disable TLS | Yes | | --ca_cert CA_CERT | CA certificate used to connect to self hosted AMaaS | Yes | diff --git a/protos/scan.proto b/protos/scan.proto index 3537f2f..db230f7 100644 --- a/protos/scan.proto +++ b/protos/scan.proto @@ -15,13 +15,14 @@ enum Stage { STAGE_INIT = 0; STAGE_RUN = 1; STAGE_FINI = 2; + STAGE_HEARTBEAT = 3; } message C2S { Stage stage = 1; string file_name = 2; uint64 rs_size = 3; - int32 offset = 4; + int64 offset = 4; bytes chunk = 5; bool trendx = 6; string file_sha1 = 7; @@ -40,10 +41,10 @@ enum Command { message S2C { Stage stage = 1; Command cmd = 2; - int32 offset = 3; + int64 offset = 3; int32 length = 4; string result = 5; - repeated int32 bulk_offset = 6; + repeated int64 bulk_offset = 6; repeated int32 bulk_length = 7; string session_id = 8; } diff --git a/tests/mock_server.py b/tests/mock_server.py index 0a604ae..0db9798 100644 --- a/tests/mock_server.py +++ b/tests/mock_server.py @@ -3,7 +3,7 @@ import grpc from amaas.grpc.protos.scan_pb2_grpc import ScanServicer -from amaas.grpc.protos.scan_pb2 import STAGE_INIT, STAGE_FINI, STAGE_RUN +from amaas.grpc.protos.scan_pb2 import STAGE_INIT, STAGE_FINI, STAGE_RUN, STAGE_HEARTBEAT from amaas.grpc.protos.scan_pb2 import CMD_QUIT, CMD_RETR from amaas.grpc.protos.scan_pb2 import S2C @@ -51,6 +51,8 @@ def getMismatchedCmdStage(self): def Run(self, request_iterator, context): count = 0 for req in request_iterator: + if req.stage == STAGE_HEARTBEAT: + continue if req.stage == STAGE_INIT: self.fsize = req.rs_size self.identifier = req.file_name diff --git a/tests/test_aio_client_sdk.py b/tests/test_aio_client_sdk.py index 4cb8916..7b739c2 100644 --- a/tests/test_aio_client_sdk.py +++ b/tests/test_aio_client_sdk.py @@ -1,3 +1,4 @@ +import asyncio import grpc import json import os @@ -5,7 +6,7 @@ import random import tempfile from concurrent import futures -from unittest.mock import patch +from unittest.mock import patch, AsyncMock, MagicMock import amaas.grpc.aio from .mock_server import MockScanServicer @@ -168,3 +169,51 @@ async def test_scan_buffer_exceptions(error_type, expected_exception): await amaas.grpc.aio.scan_buffer(handle, buffer, error_type) for cnt in range(len(expected_exception)): assert exc_info.value.args[cnt] == expected_exception[cnt] + + +# +# Testing that aio _scan_data sends heartbeat when call.read() times out +# +@pytest.mark.asyncio +async def test_aio_heartbeat(): + read_count = 0 + quit_response = amaas.grpc.scan_pb2.S2C( + cmd=amaas.grpc.scan_pb2.CMD_QUIT, + stage=amaas.grpc.scan_pb2.STAGE_FINI, + result='{"scanResult": 0}', + ) + + # Mock call.read(): simulate server not responding on first read (triggers heartbeat), + # then returning CMD_QUIT on second read to end the scan normally. + async def mock_read(): + nonlocal read_count + read_count += 1 + if read_count == 1: + await asyncio.sleep(10) # will be cancelled by wait_for timeout + return quit_response + + # Mock the gRPC call object returned by stub.Run() + mock_call = AsyncMock() + mock_call.read = mock_read + mock_call.write = AsyncMock() + mock_call.done_writing = AsyncMock() + + # Mock scan_pb2_grpc.ScanStub so no real gRPC connection is made + mock_stub = MagicMock() + mock_stub.Run = MagicMock(return_value=mock_call) + + with patch("amaas.grpc.aio.scan_pb2_grpc.ScanStub", return_value=mock_stub): + with patch("amaas.grpc.aio.heartbeat_interval_in_seconds", 0.1): + handle = MagicMock() # Mock grpc.aio.Channel + result = await amaas.grpc.aio.scan_buffer( + handle, b"test data", "test.txt", digest=False + ) + + # Verify heartbeat was sent + heartbeat_calls = [ + c + for c in mock_call.write.call_args_list + if c[0][0].stage == amaas.grpc.scan_pb2.STAGE_HEARTBEAT + ] + assert len(heartbeat_calls) == 1 + assert json.loads(result)["scanResult"] == 0 diff --git a/tests/test_client_sdk.py b/tests/test_client_sdk.py index bc23a1e..5246f6f 100644 --- a/tests/test_client_sdk.py +++ b/tests/test_client_sdk.py @@ -240,6 +240,46 @@ def test_generate_message_unknwon_stage(): assert exc_info.value.args[1] == UNKNOWN_STAGE +# +# Testing that _Pipeline.get_message returns None on timeout +# +def test_pipeline_get_message_timeout(): + pipeline = amaas.grpc._Pipeline() + result = pipeline.get_message(timeout=0.1) + assert result is None + + +# +# Testing that _generate_messages yields heartbeat when pipeline times out +# +def test_generate_message_heartbeat(): + pipeline = amaas.grpc._Pipeline() + stats = {} + f = open(TEST_DATA_FILE_NAME, "rb") + size = os.stat(TEST_DATA_FILE_NAME).st_size + + # Set initial message + server_resp = amaas.grpc.scan_pb2.C2S( + stage=amaas.grpc.scan_pb2.STAGE_INIT, + file_name=TEST_DATA_FILE_NAME, + rs_size=size, + offset=0, + chunk=None, + tags=None, + ) + pipeline.set_message(server_resp) + + with patch("amaas.grpc.heartbeat_interval_in_seconds", 0.1): + gen = amaas.grpc._generate_messages(pipeline, f, True, stats) + # First message should be the INIT message + c2s_msg = next(gen) + assert c2s_msg.stage == amaas.grpc.scan_pb2.STAGE_INIT + + # Next message should be a heartbeat (since no pipeline message is available) + c2s_msg = next(gen) + assert c2s_msg.stage == amaas.grpc.scan_pb2.STAGE_HEARTBEAT + + # # Testing the SDK scan_file method sucessfully scans a file with no virus. # @@ -250,6 +290,16 @@ def test_scan_file_success(): assert json.loads(response)["scanResult"] == 0 +def test_scan_file_identifier_is_full_path(): + full_path = os.path.join(os.path.dirname(__file__), "fake_server_cert.pem") + assert os.path.basename(full_path) != full_path + handle = grpc.insecure_channel(f"localhost:{SERVER_PORT}") + with patch("amaas.grpc._scan_data") as scan_data_mock: + amaas.grpc.scan_file(handle, full_path) + identifier = scan_data_mock.call_args.args[3] + assert identifier == full_path + + # # Testing the SDK scan_file method sucessfully failed with MSG_ID_ERR_FILE_NOT_FOUND exception. #