diff --git a/.gitignore b/.gitignore index e81f084..3e37154 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,7 @@ __pycache__/ # due to using tox and pytest .tox .cache +.ruff_cache + +# generated protobuf stubs (make proto / make build) +amaas/grpc/protos/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 6302f2e..2840eee 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,24 @@ # CHANGELOG +## 1.5.0 - 2026-09-15 + +- Add `scan_reader` (sync and aio) for scanning a data source through a reader implementing + the new `amaas.grpc.reader.AMaasReader` protocol, mirroring the Go SDK's `ScanReader` and + `AmaasClientReader`. The SDK pulls only the chunks the scan engine requests, so remote + sources (e.g. S3 objects) need not be downloaded in full. `digest` defaults to `False` + because digest calculation reads the whole data source. +- Add an S3 object scan example under `examples/scan-s3obj` demonstrating a reader backed + by ranged S3 GETs. +- Add an Azure Blob scan example under `examples/scan-azureblob` demonstrating a reader backed + by ranged blob downloads, with SAS token or `DefaultAzureCredential` auth. +- Enforce the reader contract like the Go SDK: `read_bytes` is never called beyond `data_size`, and + a short read raises the new `MSG_ID_ERR_RETRIEVE_DATA` error instead of uploading truncated data + or hashing a truncated source for `digest=True`. +- Compute sha1 and sha256 digests in a single pass, halving the reads a remote reader makes when + `digest=True`. +- `scan_file` now reports `MSG_ID_ERR_UNEXPECTED_ERROR` for `OSError` other than permission-denied; + only `PermissionError` maps to `MSG_ID_ERR_FILE_NO_PERMISSION`. + ## 1.4.8 - 2026-08-19 - Support new region ap-southeast-3 (Indonesia) diff --git a/README.md b/README.md index 2259b4e..53fbb78 100644 --- a/README.md +++ b/README.md @@ -264,6 +264,65 @@ AsyncIO Scan a file for malware and retrieves response data from the API. **_Return_** String the scanned result in JSON format. +### Scanning with a Reader + +`scan_reader` scans a data source through a reader object you provide, mirroring the Go SDK's `AmaasClientReader` interface. The SDK pulls only the chunks the scan engine requests, so a reader backed by a remote source (for example an S3 object) never needs to download the whole object. + +Implement the `amaas.grpc.reader.AMaasReader` protocol: + +```python +class AMaasReader(Protocol): + def identifier(self) -> str: + """Return the identifier of the data source, e.g. "s3://bucket/key".""" + + def data_size(self) -> int: + """Return the total size of the data source in bytes.""" + + def read_bytes(self, offset: int, length: int) -> bytes: + """Return exactly length bytes of the data source starting at offset. + The SDK never requests beyond data_size; a short read raises + MSG_ID_ERR_RETRIEVE_DATA and fails the scan.""" +``` + +Then pass the reader to `scan_reader` (or `amaas.grpc.aio.scan_reader`): + +```python +reader = MyS3ObjectReader(bucket, key) +result = amaas.grpc.scan_reader(handle, reader, tags=tags) +``` + +See [examples/scan-s3obj/scan_s3obj.py](examples/scan-s3obj/scan_s3obj.py) for a complete S3 implementation. + +#### `def amaas.grpc.scan_reader(handle: grpc.Channel, reader: AMaasReader, tags: List[str], pml: bool = False, feedback: bool = False, verbose: bool = False, digest: bool = False) -> str` + +Scan a data source through an `AMaasReader` implementation. + +**_Parameters_** + +| Parameter | Description | +| --------- | ----------------------------------------------------------------------------------------------------------- | +| handle | The grpc Channel instance was created from the init function. | +| reader | An object implementing the `AMaasReader` protocol (identifier / data_size / read_bytes). | +| tags | A list of strings to be used to tag the scan result. At most 8 tags with a maximum length of 63 characters. | +| pml | Enable PML (Predictive Machine Learning) Detection. | +| feedback | Enable SPN feedback for Predictive Machine Learning Detection | +| verbose | Enable log verbose mode | +| digest | Calculate digests for cache search and result lookup. Defaults to `False` because digest calculation reads the whole data source; with a remote reader that defeats the purpose of partial reads. Pass `True` to opt in. | + +**_Return_** +String the scanned result in JSON format. + +#### `def amaas.grpc.aio.scan_reader(handle: grpc.aio.Channel, reader: AMaasReader, tags: List[str], pml: bool = False, feedback: bool = False, verbose: bool = False, digest: bool = False) -> str` + +AsyncIO scan of a data source through an `AMaasReader` implementation. `read_bytes()` runs in a worker thread (`asyncio.to_thread`), so the event loop stays responsive while a remote reader (e.g. S3 ranged GETs) fetches a chunk. Readers stay synchronous, matching the sync client semantics. + +**_Parameters_** + +Same as `amaas.grpc.scan_reader`, with the aio Channel handle. + +**_Return_** +String the scanned result in JSON format. + ### Cleaning Up #### `def amaas.grpc.quit(handle: grpc.aio.Channel) -> None` diff --git a/VERSION b/VERSION index b2e46d1..3e1ad72 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -1.4.8 +1.5.0 \ No newline at end of file diff --git a/amaas/grpc/__init__.py b/amaas/grpc/__init__.py index 40a916f..c92d11c 100644 --- a/amaas/grpc/__init__.py +++ b/amaas/grpc/__init__.py @@ -1,29 +1,31 @@ -import threading +import io import logging -from typing import BinaryIO, List +import os +import threading +from typing import BinaryIO import grpc -import os -import io -from .protos import scan_pb2 -from .protos import scan_pb2_grpc -from .exception import AMaasException -from .exception import AMaasErrorCode -from .util import _init_by_region_util -from .util import _init_util -from .util import _validate_tags -from .util import _digest_hex -from .util import APP_NAME_HEADER, APP_NAME_FILE_SCAN +from .exception import AMaasErrorCode, AMaasException +from .protos import scan_pb2, scan_pb2_grpc +from .reader import AMaasReader, _ReaderAdapter +from .util import ( + APP_NAME_FILE_SCAN, + APP_NAME_HEADER, + _digest_hex_pair, + _init_by_region_util, + _init_util, + _validate_tags, +) logger = logging.getLogger(__name__) logger.addHandler(logging.StreamHandler()) -LOG_LEVEL = os.environ.get('LOG_LEVEL', 'INFO') +LOG_LEVEL = os.environ.get("LOG_LEVEL", "INFO") logger.setLevel(LOG_LEVEL) 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)) +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: @@ -59,7 +61,9 @@ def init(host, api_key=None, enable_tls=False, ca_cert=None): return _init_util(host, api_key, enable_tls, ca_cert, False) -def _generate_messages(pipeline: _Pipeline, data_reader: BinaryIO, bulk: bool, stats: dict) -> None: +def _generate_messages( + pipeline: _Pipeline, data_reader: BinaryIO, bulk: bool, stats: dict +) -> None: responses = [] while True: @@ -70,7 +74,14 @@ def _generate_messages(pipeline: _Pipeline, data_reader: BinaryIO, bulk: bool, s offset = r[1] length = r[2] data_reader.seek(offset) - chunk = data_reader.read(length) + try: + chunk = data_reader.read(length) + except AMaasException as err: + # The request iterator runs on a gRPC consumer thread, so + # this exception is swallowed there; stash it for the + # response loop to re-raise in the caller's thread. + stats["reader_error"] = err + raise response = scan_pb2.C2S( stage=scan_pb2.STAGE_RUN, file_name=None, @@ -80,7 +91,9 @@ def _generate_messages(pipeline: _Pipeline, data_reader: BinaryIO, bulk: bool, s ) stats["total_upload"] = stats.get("total_upload", 0) + len(chunk) else: - raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_CMD_AND_STAGE, "None", r[0]) + raise AMaasException( + AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_CMD_AND_STAGE, "None", r[0] + ) yield response responses.clear() @@ -97,7 +110,11 @@ def _generate_messages(pipeline: _Pipeline, data_reader: BinaryIO, bulk: bool, s responses.append(("INIT", message)) elif message.stage == scan_pb2.STAGE_RUN: if message.cmd != scan_pb2.CMD_RETR: - raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_CMD_AND_STAGE, message.cmd, message.stage) + raise AMaasException( + AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_CMD_AND_STAGE, + message.cmd, + message.stage, + ) length = [] offset = [] @@ -113,11 +130,17 @@ def _generate_messages(pipeline: _Pipeline, data_reader: BinaryIO, bulk: bool, s length.append(message.length) for i in range(len(length)): - logger.debug(f"stage RUN, try to read {length[i]} at offset {offset[i]}") + logger.debug( + f"stage RUN, try to read {length[i]} at offset {offset[i]}" + ) responses.append(("RUN", offset[i], length[i])) elif message.stage == scan_pb2.STAGE_FINI: if message.cmd != scan_pb2.CMD_QUIT: - raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_CMD_AND_STAGE, message.cmd, message.stage) + raise AMaasException( + AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_CMD_AND_STAGE, + message.cmd, + message.stage, + ) logger.debug("final stage, quit generating C2S messages...") break @@ -130,8 +153,17 @@ def quit(handle): handle.close() -def _scan_data(channel: grpc.Channel, data_reader: BinaryIO, size: int, identifier: str, tags: List[str], - pml: bool, feedback: bool, verbose: bool, digest: bool) -> str: +def _scan_data( + channel: grpc.Channel, + data_reader: BinaryIO, + size: int, + identifier: str, + tags: list[str], + pml: bool, + feedback: bool, + verbose: bool, + digest: bool, +) -> str: _validate_tags(tags) stub = scan_pb2_grpc.ScanStub(channel) pipeline = _Pipeline() @@ -142,27 +174,31 @@ def _scan_data(channel: grpc.Channel, data_reader: BinaryIO, size: int, identifi file_sha256 = "" if digest: - file_sha1 = "sha1:" + _digest_hex(data_reader, "sha1") - file_sha256 = "sha256:" + _digest_hex(data_reader, "sha256") + sha1_hex, sha256_hex = _digest_hex_pair(data_reader) + file_sha1 = "sha1:" + sha1_hex + file_sha256 = "sha256:" + sha256_hex try: - metadata = ( - (APP_NAME_HEADER, APP_NAME_FILE_SCAN), + metadata = ((APP_NAME_HEADER, APP_NAME_FILE_SCAN),) + responses = stub.Run( + _generate_messages(pipeline, data_reader, bulk, stats), + timeout=timeout_in_seconds, + metadata=metadata, + ) + message = scan_pb2.C2S( + stage=scan_pb2.STAGE_INIT, + file_name=identifier, + rs_size=size, + offset=0, + chunk=None, + trendx=pml, + tags=tags, + file_sha1=file_sha1, + file_sha256=file_sha256, + bulk=bulk, + spn_feedback=feedback, + verbose=verbose, ) - responses = stub.Run(_generate_messages(pipeline, data_reader, bulk, stats), timeout=timeout_in_seconds, - metadata=metadata) - message = scan_pb2.C2S(stage=scan_pb2.STAGE_INIT, - file_name=identifier, - rs_size=size, - offset=0, - chunk=None, - trendx=pml, - tags=tags, - file_sha1=file_sha1, - file_sha256=file_sha256, - bulk=bulk, - spn_feedback=feedback, - verbose=verbose) pipeline.set_message(message) @@ -176,7 +212,9 @@ def _scan_data(channel: grpc.Channel, data_reader: BinaryIO, size: int, identifi break else: logger.debug("unknown command...") - raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNKNOWN_CMD, response.cmd) + raise AMaasException( + AMaasErrorCode.MSG_ID_ERR_UNKNOWN_CMD, response.cmd + ) total_upload = stats.get("total_upload", 0) logger.debug(f"total upload {total_upload} bytes") @@ -184,20 +222,38 @@ def _scan_data(channel: grpc.Channel, data_reader: BinaryIO, size: int, identifi except AMaasException: raise except grpc.RpcError as rpc_error: + reader_error = stats.get("reader_error") + if reader_error is not None: + raise reader_error if "429" in str(rpc_error): raise AMaasException(AMaasErrorCode.MSG_ID_ERR_RATE_LIMIT_EXCEEDED) elif rpc_error.code() == grpc.StatusCode.UNAUTHENTICATED: raise AMaasException(AMaasErrorCode.MSG_ID_ERR_KEY_AUTH_FAILED) else: - raise AMaasException(AMaasErrorCode.MSG_ID_GRPC_ERROR, rpc_error.code().value[0], rpc_error.details()) + raise AMaasException( + AMaasErrorCode.MSG_ID_GRPC_ERROR, + rpc_error.code().value[0], + rpc_error.details(), + ) except Exception as err: raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_ERROR, str(err)) + reader_error = stats.get("reader_error") + if reader_error is not None: + raise reader_error + return result -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: +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 = file_name @@ -205,14 +261,62 @@ def scan_file(channel: grpc.Channel, file_name: str, tags: List[str] = None, except FileNotFoundError as err: logger.debug("File not exist: " + str(err)) raise AMaasException(AMaasErrorCode.MSG_ID_ERR_FILE_NOT_FOUND, file_name) - except (PermissionError, IOError) as err: + except PermissionError as err: # only Errno 13 reports no permission logger.debug("Permission error: " + str(err)) raise AMaasException(AMaasErrorCode.MSG_ID_ERR_FILE_NO_PERMISSION, file_name) + except OSError as err: + raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_ERROR, str(err)) return _scan_data(channel, f, n, fid, tags, pml, feedback, verbose, digest) -def scan_buffer(channel: grpc.Channel, bytes_buffer: bytes, uid: str, tags: List[str] = None, - pml: bool = False, feedback: bool = False, verbose: bool = False, digest: bool = True) -> str: +def scan_buffer( + channel: grpc.Channel, + bytes_buffer: bytes, + uid: str, + tags: list[str] = None, + pml: bool = False, + feedback: bool = False, + verbose: bool = False, + digest: bool = True, +) -> str: f = io.BytesIO(bytes_buffer) - return _scan_data(channel, f, len(bytes_buffer), uid, tags, pml, feedback, verbose, digest) + return _scan_data( + channel, f, len(bytes_buffer), uid, tags, pml, feedback, verbose, digest + ) + + +def scan_reader( + channel: grpc.Channel, + reader: AMaasReader, + tags: list[str] = None, + pml: bool = False, + feedback: bool = False, + verbose: bool = False, + digest: bool = False, +) -> str: + """Scan a data source through an AMaasReader (see amaas.grpc.reader). + + digest defaults to False because digest calculation reads the whole data + source; with a remote reader (e.g. an S3 object) that defeats the purpose + of partial reads. Pass digest=True to opt in. + """ + adapter = _ReaderAdapter(reader) + try: + size = reader.data_size() + fid = reader.identifier() + return _scan_data( + channel, + adapter, + size, + fid, + tags, + pml, + feedback, + verbose, + digest, + ) + except AMaasException: + raise + except Exception as err: + raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_ERROR, str(err)) diff --git a/amaas/grpc/aio/__init__.py b/amaas/grpc/aio/__init__.py index b491a6d..333adfe 100644 --- a/amaas/grpc/aio/__init__.py +++ b/amaas/grpc/aio/__init__.py @@ -1,29 +1,31 @@ import asyncio import io +import logging import os -from typing import BinaryIO, List +from typing import BinaryIO import grpc -import logging -from ..protos import scan_pb2 -from ..protos import scan_pb2_grpc -from ..exception import AMaasException -from ..exception import AMaasErrorCode -from ..util import _init_by_region_util -from ..util import _init_util -from ..util import _validate_tags -from ..util import _digest_hex -from ..util import APP_NAME_HEADER, APP_NAME_FILE_SCAN +from ..exception import AMaasErrorCode, AMaasException +from ..protos import scan_pb2, scan_pb2_grpc +from ..reader import AMaasReader, _ReaderAdapter +from ..util import ( + APP_NAME_FILE_SCAN, + APP_NAME_HEADER, + _digest_hex_pair, + _init_by_region_util, + _init_util, + _validate_tags, +) logger = logging.getLogger(__name__) logger.addHandler(logging.StreamHandler()) -LOG_LEVEL = os.environ.get('LOG_LEVEL', 'INFO') +LOG_LEVEL = os.environ.get("LOG_LEVEL", "INFO") logger.setLevel(LOG_LEVEL) 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)) +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): @@ -41,8 +43,17 @@ async def quit(handle): # https://github.com/grpc/grpc/blob/91083659fa88c938779dd41e57a7f97981b6c9a1/src/python/grpcio_tests/tests_aio/unit/channel_test.py#L180 -async def _scan_data(channel: grpc.Channel, data_reader: BinaryIO, size: int, identifier: str, tags: List[str], - pml: bool, feedback: bool, verbose: bool, digest: bool) -> str: +async def _scan_data( + channel: grpc.Channel, + data_reader: BinaryIO, + size: int, + identifier: str, + tags: list[str], + pml: bool, + feedback: bool, + verbose: bool, + digest: bool, +) -> str: _validate_tags(tags) stub = scan_pb2_grpc.ScanStub(channel) stats = {} @@ -52,27 +63,28 @@ async def _scan_data(channel: grpc.Channel, data_reader: BinaryIO, size: int, id file_sha256 = "" if digest: - file_sha1 = "sha1:" + _digest_hex(data_reader, "sha1") - file_sha256 = "sha256:" + _digest_hex(data_reader, "sha256") + sha1_hex, sha256_hex = await asyncio.to_thread(_digest_hex_pair, data_reader) + file_sha1 = "sha1:" + sha1_hex + file_sha256 = "sha256:" + sha256_hex try: - metadata = ( - (APP_NAME_HEADER, APP_NAME_FILE_SCAN), - ) + metadata = ((APP_NAME_HEADER, APP_NAME_FILE_SCAN),) call = stub.Run(timeout=timeout_in_seconds, metadata=metadata) - request = scan_pb2.C2S(stage=scan_pb2.STAGE_INIT, - file_name=identifier, - rs_size=size, - offset=0, - chunk=None, - tags=tags, - trendx=pml, - file_sha1=file_sha1, - file_sha256=file_sha256, - bulk=bulk, - spn_feedback=feedback, - verbose=verbose) + request = scan_pb2.C2S( + stage=scan_pb2.STAGE_INIT, + file_name=identifier, + rs_size=size, + offset=0, + chunk=None, + tags=tags, + trendx=pml, + file_sha1=file_sha1, + file_sha256=file_sha256, + bulk=bulk, + spn_feedback=feedback, + verbose=verbose, + ) await call.write(request) @@ -89,8 +101,11 @@ async def _scan_data(channel: grpc.Channel, data_reader: BinaryIO, size: int, id if response.cmd == scan_pb2.CMD_RETR: if response.stage != scan_pb2.STAGE_RUN: - raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_CMD_AND_STAGE, response.cmd, - response.stage) + raise AMaasException( + AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_CMD_AND_STAGE, + response.cmd, + response.stage, + ) length = [] offset = [] @@ -111,29 +126,37 @@ async def _scan_data(channel: grpc.Channel, data_reader: BinaryIO, size: int, id for i in range(len(length)): logger.debug(f"try to read {length[i]} at offset {offset[i]}") data_reader.seek(offset[i]) - chunk = data_reader.read(length[i]) + # to_thread keeps remote readers (e.g. S3 ranged GETs) from + # stalling the event loop; local file/buffer reads are + # equally safe off-loop. + chunk = await asyncio.to_thread(data_reader.read, length[i]) request = scan_pb2.C2S( stage=scan_pb2.STAGE_RUN, file_name=None, rs_size=0, offset=offset[i], - chunk=chunk) + chunk=chunk, + ) - stats["total_upload"] = stats.get( - "total_upload", 0) + len(chunk) + stats["total_upload"] = stats.get("total_upload", 0) + len(chunk) await call.write(request) elif response.cmd == scan_pb2.CMD_QUIT: if response.stage != scan_pb2.STAGE_FINI: - raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_CMD_AND_STAGE, response.cmd, - response.stage) + raise AMaasException( + AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_CMD_AND_STAGE, + response.cmd, + response.stage, + ) result = response.result logger.debug("receive QUIT, exit loop...") break else: logger.debug("unknown command...") - raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNKNOWN_CMD, response.cmd) + raise AMaasException( + AMaasErrorCode.MSG_ID_ERR_UNKNOWN_CMD, response.cmd + ) await call.done_writing() @@ -148,15 +171,26 @@ async def _scan_data(channel: grpc.Channel, data_reader: BinaryIO, size: int, id elif rpc_error.code() == grpc.StatusCode.UNAUTHENTICATED: raise AMaasException(AMaasErrorCode.MSG_ID_ERR_KEY_AUTH_FAILED) else: - raise AMaasException(AMaasErrorCode.MSG_ID_GRPC_ERROR, rpc_error.code().value[0], rpc_error.details()) + raise AMaasException( + AMaasErrorCode.MSG_ID_GRPC_ERROR, + rpc_error.code().value[0], + rpc_error.details(), + ) except Exception as err: raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_ERROR, str(err)) return result -async 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: +async 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 = file_name @@ -164,13 +198,65 @@ async def scan_file(channel: grpc.Channel, file_name: str, tags: List[str] = Non except FileNotFoundError as err: logger.debug("File not exist: " + str(err)) raise AMaasException(AMaasErrorCode.MSG_ID_ERR_FILE_NOT_FOUND, file_name) - except (PermissionError, IOError) as err: + except PermissionError as err: # only Errno 13 reports no permission logger.debug("Permission error: " + str(err)) raise AMaasException(AMaasErrorCode.MSG_ID_ERR_FILE_NO_PERMISSION, file_name) + except OSError as err: + raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_ERROR, str(err)) return await _scan_data(channel, f, n, fid, tags, pml, feedback, verbose, digest) -async def scan_buffer(channel: grpc.Channel, bytes_buffer: bytes, uid: str, tags: List[str] = None, - pml: bool = False, feedback: bool = False, verbose: bool = False, digest: bool = True) -> str: +async def scan_buffer( + channel: grpc.Channel, + bytes_buffer: bytes, + uid: str, + tags: list[str] = None, + pml: bool = False, + feedback: bool = False, + verbose: bool = False, + digest: bool = True, +) -> str: f = io.BytesIO(bytes_buffer) - return await _scan_data(channel, f, len(bytes_buffer), uid, tags, pml, feedback, verbose, digest) + return await _scan_data( + channel, f, len(bytes_buffer), uid, tags, pml, feedback, verbose, digest + ) + + +async def scan_reader( + channel: grpc.Channel, + reader: AMaasReader, + tags: list[str] = None, + pml: bool = False, + feedback: bool = False, + verbose: bool = False, + digest: bool = False, +) -> str: + """Scan a data source through an AMaasReader (see amaas.grpc.reader). + + read_bytes() runs in a worker thread (asyncio.to_thread), so the event + loop stays responsive while a remote reader (e.g. S3 ranged GETs) fetches + a chunk. Readers stay synchronous, matching the sync client semantics. + + digest defaults to False because digest calculation reads the whole data + source; with a remote reader (e.g. an S3 object) that defeats the purpose + of partial reads. Pass digest=True to opt in. + """ + adapter = _ReaderAdapter(reader) + try: + size = reader.data_size() + fid = reader.identifier() + return await _scan_data( + channel, + adapter, + size, + fid, + tags, + pml, + feedback, + verbose, + digest, + ) + except AMaasException: + raise + except Exception as err: + raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_ERROR, str(err)) diff --git a/amaas/grpc/exception/__init__.py b/amaas/grpc/exception/__init__.py index ceae3e9..7658dd5 100644 --- a/amaas/grpc/exception/__init__.py +++ b/amaas/grpc/exception/__init__.py @@ -20,7 +20,10 @@ def __str__(self): class AMaasErrorCode(Enum): MSG_ID_ERR_FILE_NOT_FOUND = "Failed to open file. No such file or directory %s." MSG_ID_ERR_FILE_NO_PERMISSION = "Failed to open file. Permission denied to open %s." - MSG_ID_ERR_INVALID_REGION = "%s is not a supported region, region value should be one of %s" + MSG_ID_ERR_RETRIEVE_DATA = "Attempted to read %d bytes but only retrieved %d" + MSG_ID_ERR_INVALID_REGION = ( + "%s is not a supported region, region value should be one of %s" + ) MSG_ID_ERR_MISSING_AUTH = "Must provide an API key to use the client." MSG_ID_GRPC_ERROR = "Received gRPC status code: %s, msg: %s." MSG_ID_ERR_KEY_AUTH_FAILED = "Invalid token or Api Key." @@ -28,6 +31,8 @@ class AMaasErrorCode(Enum): MSG_ID_ERR_UNKNOWN_STAGE = "Received unknown stage from server: %d" MSG_ID_ERR_UNEXPECTED_CMD_AND_STAGE = "Received unexpected command %d and stage %d." MSG_ID_ERR_UNEXPECTED_ERROR = "Unexpected error encountered. %s" - MSG_ID_ERR_RATE_LIMIT_EXCEEDED = "Raised by the SDK library to indicate http 429 too many request error." + MSG_ID_ERR_RATE_LIMIT_EXCEEDED = ( + "Raised by the SDK library to indicate http 429 too many request error." + ) MSG_ID_ERR_INVALID_TAG = "Invalid tag format: %s." MSG_ID_ERR_TAG_NUMBER_EXCEED = "Too many tags: %d." diff --git a/amaas/grpc/reader.py b/amaas/grpc/reader.py new file mode 100644 index 0000000..ae54d3e --- /dev/null +++ b/amaas/grpc/reader.py @@ -0,0 +1,73 @@ +"""Reader support for scanning data sources that are not local files. + +The AMaasReader protocol mirrors the Go SDK's AmaasClientReader interface: +the scan engine pulls the chunks it needs via read_bytes(), so a reader can +serve data from a remote source (e.g. an S3 object) without downloading it +in full. _ReaderAdapter adapts a reader to the seekable file-like object +that the scan protocol in amaas.grpc._scan_data consumes. +""" + +import io +from typing import Protocol + +from .exception import AMaasErrorCode, AMaasException + + +class AMaasReader(Protocol): + """Reader protocol consumed by amaas.grpc.scan_reader. + + Mirrors the Go SDK's AmaasClientReader (Identifier/DataSize/ReadBytes). + """ + + def identifier(self) -> str: + """Return the identifier of the data source, e.g. "s3://bucket/key".""" + ... + + def data_size(self) -> int: + """Return the total size of the data source in bytes.""" + ... + + def read_bytes(self, offset: int, length: int) -> bytes: + """Return length bytes of the data source starting at offset.""" + ... + + +class _ReaderAdapter: + """Adapt an AMaasReader to the seekable file-like object the scan + protocol expects (seek/tell/read).""" + + def __init__(self, reader: AMaasReader): + self._reader = reader + self._pos = 0 + + def seek(self, offset: int, whence: int = io.SEEK_SET) -> int: + if whence == io.SEEK_SET: + self._pos = offset + elif whence == io.SEEK_CUR: + self._pos += offset + elif whence == io.SEEK_END: + self._pos = max(0, self._reader.data_size() + offset) + else: + raise ValueError(f"invalid whence: {whence}") + if self._pos < 0: + raise ValueError("negative seek position") + return self._pos + + def tell(self) -> int: + return self._pos + + def read(self, length: int = -1) -> bytes: + size = self._reader.data_size() + if self._pos >= size: + return b"" + if length is None or length < 0 or length > size - self._pos: + length = size - self._pos + data = self._reader.read_bytes(self._pos, length) + if len(data) != length: + # A reader reporting more data_size than it serves would otherwise + # silently upload a truncated chunk or hash truncated data. + raise AMaasException( + AMaasErrorCode.MSG_ID_ERR_RETRIEVE_DATA, length, len(data) + ) + self._pos += len(data) + return data diff --git a/amaas/grpc/util.py b/amaas/grpc/util.py index b8be5df..f83c04b 100644 --- a/amaas/grpc/util.py +++ b/amaas/grpc/util.py @@ -1,8 +1,9 @@ -import grpc import hashlib -from typing import BinaryIO, List -from .exception import AMaasException -from .exception import AMaasErrorCode +from typing import BinaryIO + +import grpc + +from .exception import AMaasErrorCode, AMaasException HASH_CHUNK_SIZE = 512 * 1024 @@ -35,12 +36,45 @@ 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, 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] +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, + 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_ID_REGION] +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_ID_REGION, +] AllRegions = C1Regions + V1Regions AllValidRegions = SupportedC1Regions + SupportedV1Regions @@ -65,19 +99,21 @@ def __init__(self, key): self._key = key def __call__(self, context, callback): - callback((('authorization', self._key),), None) + callback((("authorization", self._key),), None) -def _init_util(host, api_key=None, enable_tls=False, ca_cert=None, is_aio_channel=False): +def _init_util( + host, api_key=None, enable_tls=False, ca_cert=None, is_aio_channel=False +): call_creds = None if api_key: - auth_key_str = 'ApiKey ' + api_key + auth_key_str = "ApiKey " + api_key call_creds = grpc.metadata_call_credentials(_GrpcAuth(auth_key_str)) if enable_tls: if ca_cert: # Bring Your Own Certificate case - with open(ca_cert, 'rb') as f: + with open(ca_cert, "rb") as f: ssl_creds = grpc.ssl_channel_credentials(f.read()) else: ssl_creds = grpc.ssl_channel_credentials() @@ -89,36 +125,50 @@ def _init_util(host, api_key=None, enable_tls=False, ca_cert=None, is_aio_channe else: creds = grpc.composite_channel_credentials(ssl_creds, call_creds) - channel = grpc.aio.secure_channel(host, creds) if is_aio_channel else grpc.secure_channel(host, creds) + channel = ( + grpc.aio.secure_channel(host, creds) + if is_aio_channel + else grpc.secure_channel(host, creds) + ) else: - channel = grpc.aio.insecure_channel(host) if is_aio_channel else grpc.insecure_channel(host) + channel = ( + grpc.aio.insecure_channel(host) + if is_aio_channel + else grpc.insecure_channel(host) + ) return channel -def _init_by_region_util(region, api_key, enable_tls=True, ca_cert=None, is_aio_channel=False): +def _init_by_region_util( + region, api_key, enable_tls=True, ca_cert=None, is_aio_channel=False +): mapping = { - C1_US_REGION: 'antimalware.us-1.cloudone.trendmicro.com:443', - C1_IN_REGION: 'antimalware.in-1.cloudone.trendmicro.com:443', - C1_DE_REGION: 'antimalware.de-1.cloudone.trendmicro.com:443', - C1_SG_REGION: 'antimalware.sg-1.cloudone.trendmicro.com:443', - C1_AU_REGION: 'antimalware.au-1.cloudone.trendmicro.com:443', - C1_JP_REGION: 'antimalware.jp-1.cloudone.trendmicro.com:443', - C1_GB_REGION: 'antimalware.gb-1.cloudone.trendmicro.com:443', - 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', + C1_US_REGION: "antimalware.us-1.cloudone.trendmicro.com:443", + C1_IN_REGION: "antimalware.in-1.cloudone.trendmicro.com:443", + C1_DE_REGION: "antimalware.de-1.cloudone.trendmicro.com:443", + C1_SG_REGION: "antimalware.sg-1.cloudone.trendmicro.com:443", + C1_AU_REGION: "antimalware.au-1.cloudone.trendmicro.com:443", + C1_JP_REGION: "antimalware.jp-1.cloudone.trendmicro.com:443", + C1_GB_REGION: "antimalware.gb-1.cloudone.trendmicro.com:443", + 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 if region not in SupportedV1Regions: - raise AMaasException(AMaasErrorCode.MSG_ID_ERR_INVALID_REGION, region, SupportedV1Regions) + raise AMaasException( + AMaasErrorCode.MSG_ID_ERR_INVALID_REGION, region, SupportedV1Regions + ) else: # map it to C1 region if it is V1 region c1_region = V1ToC1RegionMapping.get(region) if not c1_region: - raise AMaasException(AMaasErrorCode.MSG_ID_ERR_INVALID_REGION, region, SupportedV1Regions) + raise AMaasException( + AMaasErrorCode.MSG_ID_ERR_INVALID_REGION, region, SupportedV1Regions + ) region = c1_region host = mapping.get(region, None) @@ -127,7 +177,7 @@ def _init_by_region_util(region, api_key, enable_tls=True, ca_cert=None, is_aio_ return _init_util(host, api_key, enable_tls, ca_cert, is_aio_channel) -def _validate_tags(tags: List[str]): +def _validate_tags(tags: list[str]): if tags is not None: if len(tags) > 8: raise AMaasException(AMaasErrorCode.MSG_ID_ERR_TAG_NUMBER_EXCEED, len(tags)) @@ -137,21 +187,20 @@ def _validate_tags(tags: List[str]): raise AMaasException(AMaasErrorCode.MSG_ID_ERR_INVALID_TAG, t) -def _digest_hex(data_reader: BinaryIO, algorithm: str): - if algorithm == "sha1": - file_hash = hashlib.sha1() - elif algorithm == "sha256": - file_hash = hashlib.sha256() - else: - raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_ERROR, "unsupported hash algorithm " + algorithm) +def _digest_hex_pair(data_reader: BinaryIO) -> tuple[str, str]: + """Compute the sha1 and sha256 hex digests in a single pass over + data_reader, so a remote reader downloads its source only once.""" + sha1_hash = hashlib.sha1() + sha256_hash = hashlib.sha256() w = data_reader.tell() data_reader.seek(0) chunk = data_reader.read(HASH_CHUNK_SIZE) while chunk: - file_hash.update(chunk) + sha1_hash.update(chunk) + sha256_hash.update(chunk) chunk = data_reader.read(HASH_CHUNK_SIZE) data_reader.seek(w) - return file_hash.hexdigest() + return sha1_hash.hexdigest(), sha256_hash.hexdigest() diff --git a/examples/README.md b/examples/README.md index 6423b1a..694703e 100644 --- a/examples/README.md +++ b/examples/README.md @@ -90,6 +90,70 @@ If you plan on using a Trend Vision One region, be sure to pass in region parame python3 client_aio.py -f FILENAME -a antimalware._REGION_.cloudone.trendmicro.com:443 --tls --api_key API_KEY ``` +### Run with the S3 object scan example + +`scan-s3obj/scan_s3obj.py` scans an S3 object **in place**: it sends ranged S3 GETs for +only the chunks the scan engine requests, so the object is never downloaded in full. +The SDK example also demonstrates `amaas.grpc.scan_reader` with a custom `AMaasReader` +implementation you can copy for other data sources. + +1. Install the example's dependencies: + + ```sh + cd examples/scan-s3obj/ + python3 -m pip install -r requirements.txt + ``` + +2. Configure credentials for your S3 bucket as usual (environment, `~/.aws/credentials`, or + an instance role). Then run the example: + + ```sh + python3 scan_s3obj.py -b BUCKET -k KEY --bucketregion us-west-2 -r us-east-1 --tls --api_key API_KEY + ``` + + or with a File Security server address `-a` instead of region `-r`: + + ```sh + python3 scan_s3obj.py -b BUCKET -k KEY --bucketregion us-west-2 -a antimalware._REGION_.cloudone.trendmicro.com:443 --tls --api_key API_KEY + ``` + + Keep `--no-digest` (the default) for large objects: digest calculation reads the whole + object and would defeat the partial-read design. + +### Run with the Azure Blob scan example + +`scan-azureblob/scan_azureblob.py` scans an Azure block blob **in place**: it sends ranged +blob downloads for only the chunks the scan engine requests, so the blob is never downloaded +in full. + +1. Install the example's dependencies: + + ```sh + cd examples/scan-azureblob/ + python3 -m pip install -r requirements.txt + ``` + +2. Authenticate with one of: + + - `--sas_token`: a blob service SAS token with read permission (dev/test) + - none: `DefaultAzureCredential` (az login, managed identity, environment variables). + Your account needs the `Storage Blob Data Reader` role on the storage account. + +3. Run the example: + + ```sh + python3 scan_azureblob.py -u https://ACCOUNT.blob.core.windows.net -c CONTAINER -b BLOB --sas_token "$SAS" -r us-east-1 --tls --api_key API_KEY + ``` + + or with a File Security server address `-a` instead of region `-r`: + + ```sh + python3 scan_azureblob.py -u https://ACCOUNT.blob.core.windows.net -c CONTAINER -b BLOB -a antimalware._REGION_.cloudone.trendmicro.com:443 --tls --api_key API_KEY + ``` + + Keep `--no-digest` (the default) for large blobs: digest calculation reads the whole + blob and would defeat the partial-read design. + ## File Security Post Scan Actions Actions to perform after scanning files with Trend Vision Oneā„¢ File Security diff --git a/examples/scan-azureblob/requirements.txt b/examples/scan-azureblob/requirements.txt new file mode 100644 index 0000000..6bcfd20 --- /dev/null +++ b/examples/scan-azureblob/requirements.txt @@ -0,0 +1,2 @@ +azure-storage-blob>=12.20 +azure-identity>=1.16 diff --git a/examples/scan-azureblob/scan_azureblob.py b/examples/scan-azureblob/scan_azureblob.py new file mode 100644 index 0000000..49f1cb1 --- /dev/null +++ b/examples/scan-azureblob/scan_azureblob.py @@ -0,0 +1,133 @@ +"""Example: scan an Azure Blob Storage blob without downloading it locally. + +Implements the same reader contract as the scan-s3obj example: +get_blob_properties for the blob size and ranged download_blob calls for +partial reads, so only the chunks the scan engine asks for cross the wire. + +Authentication: pass --sas_token for SAS token auth (dev/test), otherwise +DefaultAzureCredential is used (az login, managed identity, environment). + +The reader satisfies amaas.grpc.scan_reader's reader protocol: + identifier() -> scan result fileName, e.g. "azure://account/container/blob" + data_size() -> total blob size in bytes + read_bytes(offset, length) -> bytes for the requested range +""" + +import argparse +import sys +import time + +import amaas.grpc +from azure.core.exceptions import HttpResponseError +from azure.identity import DefaultAzureCredential +from azure.storage.blob import BlobClient + + +class AzureBlobReader: + """AmaasClientReader equivalent backed by an Azure block blob.""" + + def __init__( + self, account_url: str, container: str, blob: str, sas_token: str = None + ): + credential = sas_token.lstrip("?") if sas_token else DefaultAzureCredential() + self._client = BlobClient( + account_url=account_url, + container_name=container, + blob_name=blob, + credential=credential, + ) + self._host = account_url.split("//", 1)[-1].rstrip("/") + self._container = container + self._blob = blob + self._size = self._get_blob_size() + + def _get_blob_size(self) -> int: + size = self._client.get_blob_properties().size + if size is None: + raise RuntimeError("unable to get blob size from Azure Storage") + return size + + def identifier(self) -> str: + return f"azure://{self._host}/{self._container}/{self._blob}" + + def data_size(self) -> int: + return self._size + + def read_bytes(self, offset: int, length: int) -> bytes: + return self._client.download_blob(offset=offset, length=length).readall() + + +def main() -> int: + parser = argparse.ArgumentParser(description="Scan an Azure blob in place") + parser.add_argument( + "-u", + "--account_url", + required=True, + help="storage account URL, e.g. https://ACCOUNT.blob.core.windows.net", + ) + parser.add_argument("-c", "--container", required=True, help="blob container name") + parser.add_argument("-b", "--blob", required=True, help="blob name") + parser.add_argument( + "--sas_token", help="blob service SAS token; omit to use DefaultAzureCredential" + ) + parser.add_argument("-r", "--region", help="File Security region; e.g. us-east-1") + parser.add_argument("-a", "--addr", help="gRPC server address (self hosted)") + parser.add_argument("--api_key", help="api key for authentication") + parser.add_argument("--tls", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("--ca_cert", help="CA certificate for self hosted AMaaS server") + parser.add_argument("--pml", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument( + "--feedback", action=argparse.BooleanOptionalAction, default=False + ) + parser.add_argument( + "-v", "--verbose", action=argparse.BooleanOptionalAction, default=False + ) + parser.add_argument("-t", "--tags", nargs="+", help="list of tags") + parser.add_argument( + "--digest", + action=argparse.BooleanOptionalAction, + default=False, + help="calculate digests; downloads the whole blob, so keep it off for large blobs", + ) + args = parser.parse_args() + + if args.region: + handle = amaas.grpc.init_by_region( + args.region, args.api_key, args.tls, args.ca_cert + ) + elif args.addr: + handle = amaas.grpc.init(args.addr, args.api_key, args.tls, args.ca_cert) + else: + parser.error("either -r/--region or -a/--addr is required") + return 2 + + try: + reader = AzureBlobReader( + args.account_url, args.container, args.blob, args.sas_token + ) + s = time.perf_counter() + result = amaas.grpc.scan_reader( + handle, + reader, + tags=args.tags, + pml=args.pml, + feedback=args.feedback, + verbose=args.verbose, + digest=args.digest, + ) + elapsed = time.perf_counter() - s + print(f"scan executed in {elapsed:0.2f} seconds.") + print(result) + return 0 + except HttpResponseError as err: + print(f"Azure Storage error: {err}") + return 1 + except Exception as err: # noqa: BLE001 + print(err) + return 1 + finally: + amaas.grpc.quit(handle) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/scan-s3obj/requirements.txt b/examples/scan-s3obj/requirements.txt new file mode 100644 index 0000000..b2821d5 --- /dev/null +++ b/examples/scan-s3obj/requirements.txt @@ -0,0 +1 @@ +boto3>=1.34 diff --git a/examples/scan-s3obj/scan_s3obj.py b/examples/scan-s3obj/scan_s3obj.py new file mode 100644 index 0000000..081579d --- /dev/null +++ b/examples/scan-s3obj/scan_s3obj.py @@ -0,0 +1,110 @@ +"""Example: scan an AWS S3 object without downloading it locally. + +Implements the same reader contract as the Go SDK's scan-s3obj example: +GetObjectAttributes for the object size and ranged GetObject calls for +partial reads, so only the chunks the scan engine asks for cross the wire. + +The reader satisfies amaas.grpc.scan_reader's reader protocol: + identifier() -> scan result fileName, e.g. "s3://bucket/key" + data_size() -> total object size in bytes + read_bytes(offset, length) -> bytes for the requested range +""" + +import argparse +import sys +import time + +import boto3 +from botocore.exceptions import ClientError + +import amaas.grpc + + +class S3ObjectReader: + """AmaasClientReader equivalent backed by an S3 object.""" + + def __init__(self, bucket: str, key: str, region: str): + self._s3 = boto3.client("s3", region_name=region) + self._bucket = bucket + self._key = key + self._size = self._get_object_size() + + def _get_object_size(self) -> int: + attr = self._s3.get_object_attributes( + Bucket=self._bucket, Key=self._key, ObjectAttributes=["ObjectSize"] + ) + size = attr.get("ObjectSize") + if size is None: + raise RuntimeError("unable to get object size from S3") + return size + + def identifier(self) -> str: + return f"s3://{self._bucket}/{self._key}" + + def data_size(self) -> int: + return self._size + + def read_bytes(self, offset: int, length: int) -> bytes: + rng = f"bytes={offset}-{offset + length - 1}" + resp = self._s3.get_object(Bucket=self._bucket, Key=self._key, Range=rng) + return resp["Body"].read() + + +def main() -> int: + parser = argparse.ArgumentParser(description="Scan an S3 object in place") + parser.add_argument("-b", "--bucket", required=True, help="S3 bucket name") + parser.add_argument("-k", "--key", required=True, help="S3 object key") + parser.add_argument("--bucketregion", default="us-west-2", help="region of the S3 bucket") + parser.add_argument("-r", "--region", help="File Security region; e.g. us-east-1") + parser.add_argument("-a", "--addr", help="gRPC server address (self hosted)") + parser.add_argument("--api_key", help="api key for authentication") + parser.add_argument("--tls", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("--ca_cert", help="CA certificate for self hosted AMaaS server") + parser.add_argument("--pml", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("--feedback", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("-v", "--verbose", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("-t", "--tags", nargs="+", help="list of tags") + parser.add_argument( + "--digest", + action=argparse.BooleanOptionalAction, + default=False, + help="calculate digests; downloads the whole object, so keep it off for large objects", + ) + args = parser.parse_args() + + if args.region: + handle = amaas.grpc.init_by_region(args.region, args.api_key, args.tls, args.ca_cert) + elif args.addr: + handle = amaas.grpc.init(args.addr, args.api_key, args.tls, args.ca_cert) + else: + parser.error("either -r/--region or -a/--addr is required") + return 2 + + try: + reader = S3ObjectReader(args.bucket, args.key, args.bucketregion) + s = time.perf_counter() + result = amaas.grpc.scan_reader( + handle, + reader, + tags=args.tags, + pml=args.pml, + feedback=args.feedback, + verbose=args.verbose, + digest=args.digest, + ) + elapsed = time.perf_counter() - s + print(f"scan executed in {elapsed:0.2f} seconds.") + print(result) + return 0 + except ClientError as err: + print(f"S3 error: {err}") + return 1 + except Exception as err: # noqa: BLE001 + print(err) + return 1 + finally: + amaas.grpc.quit(handle) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/ruff.toml b/ruff.toml new file mode 100644 index 0000000..7ffdd0d --- /dev/null +++ b/ruff.toml @@ -0,0 +1,14 @@ +# Ruff replaces flake8 with the same rule set (pycodestyle E/W + pyflakes F). +# Kept in ruff.toml instead of pyproject.toml because pyproject.toml holds a +# `version = _VERSION_` placeholder until `make version` fills it in, which +# makes it invalid TOML for tools that run before the fill-in (e.g. tox lint). +target-version = "py39" +line-length = 120 + +[lint] +select = ["E", "W", "F"] +ignore = [ + "E203", # whitespace before : + "E741", # ambiguous variable name + "E501", # line too long +] diff --git a/tests/test_aio_client_sdk.py b/tests/test_aio_client_sdk.py index 7b739c2..08482cf 100644 --- a/tests/test_aio_client_sdk.py +++ b/tests/test_aio_client_sdk.py @@ -1,18 +1,17 @@ import asyncio -import grpc import json import os -import pytest import random import tempfile from concurrent import futures -from unittest.mock import patch, AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import amaas.grpc.aio -from .mock_server import MockScanServicer -from amaas.grpc.exception import AMaasErrorCode -from amaas.grpc.exception import AMaasException +import grpc +import pytest +from amaas.grpc.exception import AMaasErrorCode, AMaasException +from .mock_server import MockScanServicer NUM_DATA_LOOP = 128 _, TEST_DATA_FILE_NAME = tempfile.mkstemp() @@ -106,6 +105,19 @@ async def test_scan_file_no_permission(): os.remove(NOT_PERMISSION_FILE) +# +# Testing the SDK scan_file method failed with MSG_ID_ERR_UNEXPECTED_ERROR +# for an OSError that is neither not-found nor permission-denied. +# +@pytest.mark.asyncio +async def test_scan_file_other_os_error(): + handle = grpc.aio.insecure_channel(f"localhost:{SERVER_PORT}") + dir_path = os.path.dirname(os.path.realpath(__file__)) + with pytest.raises(AMaasException) as exc_info: + await amaas.grpc.scan_file(handle, dir_path) + assert exc_info.value.args[0] == AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_ERROR + + # # Testing the SDK scan_buffer method succeeded without viruses. # diff --git a/tests/test_client_sdk.py b/tests/test_client_sdk.py index 5246f6f..ae27ead 100644 --- a/tests/test_client_sdk.py +++ b/tests/test_client_sdk.py @@ -1,7 +1,5 @@ -import grpc import json import os -import pytest import random import tempfile import uuid @@ -9,10 +7,11 @@ from unittest.mock import patch import amaas.grpc -from .mock_server import MockScanServicer -from amaas.grpc.exception import AMaasErrorCode -from amaas.grpc.exception import AMaasException +import grpc +import pytest +from amaas.grpc.exception import AMaasErrorCode, AMaasException +from .mock_server import MockScanServicer NUM_DATA_LOOP = 128 _, TEST_DATA_FILE_NAME = tempfile.mkstemp() @@ -305,7 +304,7 @@ def test_scan_file_identifier_is_full_path(): # def test_scan_file_not_found(): handle = grpc.insecure_channel(f"localhost:{SERVER_PORT}") - NOT_EXIST_FILE = f"{str(uuid.uuid4())}.txt" + NOT_EXIST_FILE = f"{uuid.uuid4()!s}.txt" with pytest.raises(AMaasException) as exc_info: amaas.grpc.scan_file(handle, NOT_EXIST_FILE) assert exc_info.value.args[0] == AMaasErrorCode.MSG_ID_ERR_FILE_NOT_FOUND @@ -327,6 +326,18 @@ def test_scan_file_no_permission(): os.remove(NOT_PERMISSION_FILE) +# +# Testing the SDK scan_file method failed with MSG_ID_ERR_UNEXPECTED_ERROR +# for an OSError that is neither not-found nor permission-denied. +# +def test_scan_file_other_os_error(): + handle = grpc.insecure_channel(f"localhost:{SERVER_PORT}") + dir_path = os.path.dirname(os.path.realpath(__file__)) + with pytest.raises(AMaasException) as exc_info: + amaas.grpc.scan_file(handle, dir_path) + assert exc_info.value.args[0] == AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_ERROR + + # # Testing the SDK scan_buffer method sucessfully scans a file with no virus. # diff --git a/tests/test_reader_sdk.py b/tests/test_reader_sdk.py new file mode 100644 index 0000000..1820712 --- /dev/null +++ b/tests/test_reader_sdk.py @@ -0,0 +1,220 @@ +"""Unit tests for the AMaasReader adapter and the scan_reader APIs. + +Scans run against MockScanServicer (see mock_server.py) exactly like the +scan_buffer tests: the mock reads random ranges from the data source, which +exercises _ReaderAdapter.seek/read for every chunk the "engine" requests. +""" + +import hashlib +import json +from concurrent import futures + +import amaas.grpc +import amaas.grpc.aio +import grpc +import pytest +from amaas.grpc.reader import _ReaderAdapter + +from .mock_server import MockScanServicer + +NUM_DATA_LOOP = 128 + + +class BytesBufferReader: + """AMaasReader implementation over an in-memory buffer, mirroring the + Go SDK's AmaasClientBufferReader.""" + + def __init__(self, buffer: bytes, identifier: str): + self._buffer = buffer + self._identifier = identifier + + def identifier(self) -> str: + return self._identifier + + def data_size(self) -> int: + return len(self._buffer) + + def read_bytes(self, offset: int, length: int) -> bytes: + return self._buffer[offset : offset + length] + + +@pytest.fixture(scope="module") +def grpc_server(): + server = grpc.server(futures.ThreadPoolExecutor(max_workers=10)) + + amaas.grpc.protos.scan_pb2_grpc.add_ScanServicer_to_server( + MockScanServicer(), server + ) + # Port 0 lets the OS pick a free port; a fixed random port can collide + # with an in-use one and fail the whole module. + server_port = server.add_insecure_port("[::]:0") + server.start() + yield server_port + server.stop(None) + + +@pytest.fixture(scope="session") +def test_data(): + # Same deterministic pattern as the scan_file/scan_buffer tests: + # NUM_DATA_LOOP repetitions of bytes 0x00..0xff. + data = bytearray() + for _ in range(NUM_DATA_LOOP): + for i in range(256): + data += i.to_bytes(1, byteorder="big") + return bytes(data) + + +# +# _ReaderAdapter behaviour over an AMaasReader. +# +def test_reader_adapter_seek_tell_read(test_data): + reader = BytesBufferReader(test_data, "adapter-test") + adapter = _ReaderAdapter(reader) + + assert adapter.tell() == 0 + + # Seek + read a chunk; content must match the source bytes. + adapter.seek(300) + chunk = adapter.read(100) + assert chunk == test_data[300:400] + assert adapter.tell() == 400 + + # Reading past EOF returns b"" instead of raising. + adapter.seek(reader.data_size()) + assert adapter.read(10) == b"" + + +def test_reader_adapter_read_clamps_to_size(test_data): + reader = BytesBufferReader(test_data, "adapter-test") + adapter = _ReaderAdapter(reader) + + # A read that extends past the end returns only the remaining bytes. + adapter.seek(reader.data_size() - 10) + assert adapter.read(100) == test_data[-10:] + + +def test_reader_adapter_seek_whence(test_data): + reader = BytesBufferReader(test_data, "adapter-test") + adapter = _ReaderAdapter(reader) + + import io + + adapter.seek(100) + adapter.seek(50, io.SEEK_CUR) + assert adapter.tell() == 150 + adapter.seek(-40, io.SEEK_END) + assert adapter.tell() == reader.data_size() - 40 + + +def test_reader_adapter_digest_matches_direct_hash(test_data): + # The digest path reads the source through the adapter once and returns + # both hashes; they must equal hashing the buffer directly. + from amaas.grpc import util + + reader = BytesBufferReader(test_data, "digest-test") + adapter = _ReaderAdapter(reader) + + sha1 = hashlib.sha1(test_data).hexdigest() + sha256 = hashlib.sha256(test_data).hexdigest() + assert util._digest_hex_pair(adapter) == (sha1, sha256) + + +class ShortReadReader: + """AMaasReader that reports the full data_size but serves fewer bytes + than requested, e.g. a remote source truncating ranged reads.""" + + def __init__(self, buffer: bytes, served: int): + self._buffer = buffer + self._served = served + + def identifier(self) -> str: + return "short-read" + + def data_size(self) -> int: + return len(self._buffer) + + def read_bytes(self, offset: int, length: int) -> bytes: + return self._buffer[offset : offset + self._served] + + +def test_reader_adapter_short_read_raises(test_data): + from amaas.grpc.exception import AMaasErrorCode, AMaasException + + reader = ShortReadReader(test_data, served=2) + adapter = _ReaderAdapter(reader) + + adapter.seek(0) + with pytest.raises(AMaasException) as exc_info: + adapter.read(100) + assert exc_info.value.args[0] == AMaasErrorCode.MSG_ID_ERR_RETRIEVE_DATA + assert exc_info.value.args[1] == 100 + assert exc_info.value.args[2] == 2 + + +# +# scan_reader against the mock scan server. +# +def test_scan_reader_success(test_data, grpc_server): + handle = grpc.insecure_channel(f"localhost:{grpc_server}") + reader = BytesBufferReader(test_data, "good_sample.zip") + response = amaas.grpc.scan_reader(handle, reader) + assert json.loads(response)["scanResult"] == 0 + + +def test_scan_reader_virus(test_data, grpc_server): + handle = grpc.insecure_channel(f"localhost:{grpc_server}") + reader = BytesBufferReader(test_data, MockScanServicer.IDENTIFIER_VIRUS) + response = amaas.grpc.scan_reader(handle, reader) + jobj = json.loads(response) + assert jobj["scanResult"] == 1 + assert jobj["foundMalwares"] == ["virus1", "virus2"] + + +def test_scan_reader_digest(test_data, grpc_server): + # digest=True forces the SDK to read the whole source for hashing; the + # scan still completes and the mock's chunks keep matching. + handle = grpc.insecure_channel(f"localhost:{grpc_server}") + reader = BytesBufferReader(test_data, "digest-test") + response = amaas.grpc.scan_reader(handle, reader, digest=True) + assert json.loads(response)["scanResult"] == 0 + + +def test_scan_reader_short_read_raises(test_data, grpc_server): + from amaas.grpc.exception import AMaasErrorCode, AMaasException + + # The scan engine keeps requesting chunks; a reader serving fewer bytes + # than requested must fail the scan instead of uploading truncated data. + handle = grpc.insecure_channel(f"localhost:{grpc_server}") + reader = ShortReadReader(test_data, served=2) + with pytest.raises(AMaasException) as exc_info: + amaas.grpc.scan_reader(handle, reader) + assert exc_info.value.args[0] == AMaasErrorCode.MSG_ID_ERR_RETRIEVE_DATA + + +@pytest.mark.asyncio +async def test_scan_reader_short_read_raises_aio(test_data, grpc_server): + from amaas.grpc.exception import AMaasErrorCode, AMaasException + + handle = grpc.aio.insecure_channel(f"localhost:{grpc_server}") + reader = ShortReadReader(test_data, served=2) + with pytest.raises(AMaasException) as exc_info: + await amaas.grpc.aio.scan_reader(handle, reader) + assert exc_info.value.args[0] == AMaasErrorCode.MSG_ID_ERR_RETRIEVE_DATA + + +@pytest.mark.asyncio +async def test_scan_reader_success_aio(test_data, grpc_server): + handle = grpc.aio.insecure_channel(f"localhost:{grpc_server}") + reader = BytesBufferReader(test_data, "good_sample.zip") + response = await amaas.grpc.aio.scan_reader(handle, reader) + assert json.loads(response)["scanResult"] == 0 + + +@pytest.mark.asyncio +async def test_scan_reader_virus_aio(test_data, grpc_server): + handle = grpc.aio.insecure_channel(f"localhost:{grpc_server}") + reader = BytesBufferReader(test_data, MockScanServicer.IDENTIFIER_VIRUS) + response = await amaas.grpc.aio.scan_reader(handle, reader) + jobj = json.loads(response) + assert jobj["scanResult"] == 1 + assert jobj["foundMalwares"] == ["virus1", "virus2"] diff --git a/tox.ini b/tox.ini index 2568d88..f0730a8 100644 --- a/tox.ini +++ b/tox.ini @@ -35,7 +35,7 @@ deps = check-manifest >= 0.42 # If your project uses README.rst, uncomment the following: # readme_renderer - flake8 + ruff pytest build twine @@ -45,15 +45,9 @@ allowlist_externals= commands = make -f {toxinidir}/Makefile clean - flake8 . + ruff check . # check-manifest requires setup.py, skipping for pyproject.toml projects # check-manifest --ignore 'tox.ini,tests/**,examples/**,docs/**,*.md,Pipfile*,**/scan.proto,venv/**,.venv/**' make -f {toxinidir}/Makefile build # pytest tests {posargs} python -m twine check dist/* - -[flake8] -exclude = .tox,*.egg,build,data,venv,.venv,env,.env -select = E,W,F -max-line-length = 120 -extend-ignore = E203, E741, E501