update to latest version: v1.1.0 (#16)

.github/workflows/publish-to-pypi.yml

@@ -9,7 +9,8 @@
 name: Publish Python 🐍 distributions 📦 to PyPI and TestPyPI
 
 on:
-  push:
+  release:
+    types: [published]
 
 permissions:
   contents: read
@@ -34,7 +35,7 @@ jobs:
 
     - name: Build a binary wheel and a source tarball
       run: >-
-        make -f ./tools/Makefile build
+        make build
 
     - name: tox
       run: >-

.github/workflows/unit-test.yml

@@ -0,0 +1,26 @@
+name: unittest
+on:
+  pull_request:
+    branches:
+      - "main"
+
+jobs:
+  unit-test:
+    runs-on: ubuntu-latest
+    steps:
+      # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it
+      - uses: actions/checkout@v3
+
+      # Set up python 3.9
+      - name: Prepare python env
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.9"
+
+      # Install pipenv
+      - name: Install pipenv
+        run: pip install pipenv
+
+      # Pack and publish Python SDK
+      - name: Exec py client unit test
+        run: make test

.gitignore

@@ -0,0 +1,13 @@
+# general things to ignore
+build/
+dist/
+*.egg-info/
+*.egg
+*.py[cod]
+__pycache__/
+*.so
+*~
+
+# due to using tox and pytest
+.tox
+.cache

CHANGELOG.md

@@ -1,16 +1,33 @@
 # CHANGELOG
 
+## 1.1.0 - 2024-04-03
+
+* Update protos
+* Enable PML (Predictive Machine Learning) detection and smart feedback
+* Enable bulk mode
+* Enable India region
+* Support for scanning large files (over 2GB)
+
+## 1.0.5 - 2023-12-28
+
+* fix linting issues
+
 ## 1.0.4 - 2023-05-18
+
 * set default timeout_in_seconds to 180 seconds
 
 ## 1.0.3 - 2023-05-10
+
 * Change LICENSE
 
 ## 1.0.2 - 2023-05-10
+
 * Change README.md
 
 ## 1.0.1 - 2023-05-04
+
 * Add scan_buffer() function
-  
+
 ## 1.0.0 - 2023-05-01
-* Initial release
+
+* Initial release

tools/Makefile → Makefile

@@ -1,22 +1,24 @@
-# ROOT_DIR := $(abspath ../../)
-# include $(ROOT_DIR)/build-scripts/variables.mk
-
 PIPY_URL ?= https://upload.pypi.org/legacy/
 TOKEN ?=
+VERSION := $(shell cat VERSION | tr -d '\n')
 
-build: clean
+proto:
 	pipenv sync --dev
-	pipenv run python -m grpc_tools.protoc -Iamaas/grpc/protos=./amaas/protos \
+	pipenv run python -m grpc_tools.protoc -Iamaas/grpc/protos=./protos \
 		--python_out=. \
 		--pyi_out=. \
 		--grpc_python_out=. \
-		./amaas/protos/scan.proto
+		./protos/scan.proto
+
+build: proto
 	pipenv run pipenv-setup sync
 	pipenv run python setup.py sdist bdist_wheel
 
+test: proto
+	pipenv run pytest tests
+
 upload:
-	#pipenv run twine upload ./dist/*.whl
 	pipenv run twine upload --repository-url $(PIPY_URL) -u __token__ -p $(TOKEN) ./dist/*.whl --skip-existing
 
 clean:
-	@rm -rf dist build amaas/*.egg-info amaas/grpc/protos/
+	@rm -rf dist build *.egg-info amaas/grpc/protos/

NOTICE

@@ -1,2 +1,2 @@
 Cloud One VSAPI Python Client
-Copyright 2023 Trend Micro Inc.
+Copyright 2023 Trend Micro Inc.

Pipfile

@@ -14,7 +14,8 @@ pytest = "~=7.1"
 twine = "~=4.0"
 pipenv-setup = "~=3.2"
 vistir = "==0.6.1"
-autopep8 = "~=2.0"
+pytest-mock = "~=3.11"
+pytest-asyncio = "~=0.21"
 
 [requires]
 python_version = "3"

README.md

@@ -20,7 +20,7 @@ Install the File Security SDK package with pip:
 
 ## Obtain an API Key
 
-The File Security SDK requires a valid API Key provided as parameter to the SDK client object. It can accept Trend Vision One API keys. 
+The File Security SDK requires a valid API Key provided as parameter to the SDK client object. It can accept Trend Vision One API keys.
 
 When obtaining the API Key, ensure that the API Key is associated with the region that you plan to use. It is important to note that Trend Vision One API Keys are associated with different regions, please refer to the region flag below to obtain a better understanding of the valid regions associated with the respective API Key.
 
@@ -29,9 +29,9 @@ If you plan on using a Trend Vision One region, be sure to pass in region parame
 1. Login to the Trend Vision One.
 2. Create a new Trend Vision One API key:
 
-* Navigate to the Trend Vision One User Roles page.
-* Verify that there is a role with the "Run file scan via SDK" permissions enabled. If not, create a role by clicking on "Add Role" and "Save" once finished.
-* Directly configure a new key on the Trend Vision One API Keys page, using the role which contains the "Run file scan via SDK" permission. It is advised to set an expiry time for the API key and make a record of it for future reference.
+- Navigate to the Trend Vision One User Roles page.
+- Verify that there is a role with the "Run file scan via SDK" permissions enabled. If not, create a role by clicking on "Add Role" and "Save" once finished.
+- Directly configure a new key on the Trend Vision One API Keys page, using the role which contains the "Run file scan via SDK" permission. It is advised to set an expiry time for the API key and make a record of it for future reference.
 
 ## Run SDK
 
@@ -52,12 +52,14 @@ If you plan on using a Trend Vision One region, be sure to pass in region parame
 
 3. Current Python examples support following command line arguments
 
-   | Command Line Arguments                 | Value                    | Optional |
-   | :------------------ | :----------------------- | :------- |
-   | --region or -r | 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-southeast-1`, `ap-southeast-2`, `ap-northeast-1` | Yes, either -r or -a |
-   | --addr or -a   | Trend Vision One File Security server, such as: antimalware.__REGION__.cloudone.trendmicro.com:443 | Yes, either -r or -a      |
-   | --api_key      | Vision One API Key              | No       |
-   | --filename or -f |        File to be scanned            | No       |
+   | Command Line Arguments | Value                                                                                                                                                                               | Optional             |
+   |------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------|
+   | --region or -r         | 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-southeast-1`, `ap-southeast-2`, `ap-northeast-1` | Yes, either -r or -a |
+   | --addr or -a           | Trend Vision One File Security server, such as: antimalware.__REGION__.cloudone.trendmicro.com:443                                                                                  | Yes, either -r or -a |
+   | --api_key              | Vision One API Key                                                                                                                                                                  | No                   |
+   | --filename or -f       | File to be scanned                                                                                                                                                                  | No                   |
+   | --pml                  | Predictive Machine Learning                                                                                                                                                         | Yes                  |
+   | --tags or -f           | List of tags                                                                                                                                                                        | Yes                  |
 
 4. Run one of the examples.
 
@@ -82,46 +84,3 @@ If you plan on using a Trend Vision One region, be sure to pass in region parame
    ```sh
    python3 client_aio.py -f FILENAME -a antimalware._REGION_.cloudone.trendmicro.com:443 --tls true --api_key API_KEY
    ```
-
-### Code Examples
-
-```python
-import json
-import amaas.grpc
-
-handle = amaas.grpc.init(YOUR_FILE_SECURITY_SERVER, YOUR_VISION_ONE_KEY, True)
-
-result = amaas.grpc.scan_file(args.filename, handle)
-print(result)
-
-result_json = json.loads(result)
-print("Got scan result: %d" % result_json['scanResult'])
-
-amaas.grpc.quit(handle)
-
-```
-
-to use asyncio with  coroutines and tasks,
-
-```python:
-import json
-import pprint
-import asyncio
-import amaas.grpc.aio
-
-async def scan_files():
-   handle = amaas.grpc.aio.init(YOUR_FILE_SECURITY_SERVER, YOUR_VISION_ONE_KEY, True)
-
-   tasks = [asyncio.create_task(amaas.grpc.aio.scan_file(file_name, handle))]
-
-   scan_results = await asyncio.gather(*tasks)
-
-   for scan_result in scan_results:
-      pprint.pprint(json.loads(scan_result))
-
-   await amaas.grpc.aio.quit(handle)
-
-
-asyncio.run(scan_files())
-
-```

VERSION

@@ -0,0 +1 @@
+1.1.0

amaas/grpc/__init__.py

@@ -56,33 +56,56 @@ 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, stats: dict) -> None:
+def _generate_messages(pipeline: _Pipeline, data_reader: BinaryIO, bulk: bool, stats: dict) -> None:
+    responses = []
+
     while True:
+        for r in responses:
+            if r[0] == "INIT":
+                response = r[1]
+            elif r[0] == "RUN":
+                offset = r[1]
+                length = r[2]
+                data_reader.seek(offset)
+                chunk = data_reader.read(length)
+                response = scan_pb2.C2S(
+                    stage=scan_pb2.STAGE_RUN,
+                    file_name=None,
+                    rs_size=0,
+                    offset=offset,
+                    chunk=chunk,
+                )
+                stats["total_upload"] = stats.get("total_upload", 0) + len(chunk)
+            else:
+                raise AMaasException(AMaasErrorCode.MSG_ID_ERR_UNEXPECTED_CMD_AND_STAGE, "None", r[0])
+            yield response
+
+        responses.clear()
         message = pipeline.get_message()
 
         if message.stage == scan_pb2.STAGE_INIT:
             logger.debug("stage INIT")
-            yield message
+            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)
 
-            logger.debug(
-                f"stage RUN, try to read {message.length} at offset {message.offset}")
+            length = []
+            offset = []
 
-            data_reader.seek(message.offset)
-            chunk = data_reader.read(message.length)
+            if bulk:
+                offset = message.bulk_offset[:]
+                length = message.bulk_length[:]
 
-            message = scan_pb2.C2S(
-                stage=scan_pb2.STAGE_RUN,
-                file_name=None,
-                rs_size=0,
-                offset=data_reader.tell(),
-                chunk=chunk)
+                if len(length) > 1:
+                    logger.debug("bulk transfer triggered")
+            else:
+                offset.append(message.offset)
+                length.append(message.length)
 
-            stats["total_upload"] = stats.get(
-                "total_upload", 0) + len(chunk)
-            yield message
+            for i in range(len(length)):
+                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)
@@ -99,27 +122,31 @@ def quit(handle):
 
 
 def _scan_data(channel: grpc.Channel, data_reader: BinaryIO, size: int, identifier: str, tags: List[str],
-               app_name: str) -> str:
+               pml: bool, feedback: bool) -> str:
     _validate_tags(tags)
     stub = scan_pb2_grpc.ScanStub(channel)
     pipeline = _Pipeline()
     stats = {}
     result = None
+    bulk = True
 
     try:
         metadata = (
-            (APP_NAME_HEADER, app_name),
+            (APP_NAME_HEADER, APP_NAME_FILE_SCAN),
         )
-        responses = stub.Run(_generate_messages(pipeline, data_reader, stats), timeout=timeout_in_seconds,
+        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="sha1:" + _digest_hex(data_reader, "sha1"),
-                               file_sha256="sha256:" + _digest_hex(data_reader, "sha256"))
+                               file_sha256="sha256:" + _digest_hex(data_reader, "sha256"),
+                               bulk=bulk,
+                               spn_feedback=feedback)
 
         pipeline.set_message(message)
 
@@ -129,7 +156,7 @@ def _scan_data(channel: grpc.Channel, data_reader: BinaryIO, size: int, identifi
             elif response.cmd == scan_pb2.CMD_QUIT:
                 result = response.result
                 pipeline.set_message(response)
-                logger.debug("receive QUIT, exit loop...\n")
+                logger.debug("receive QUIT, exit loop...")
                 break
             else:
                 logger.debug("unknown command...")
@@ -153,7 +180,8 @@ def _scan_data(channel: grpc.Channel, data_reader: BinaryIO, size: int, identifi
     return result
 
 
-def scan_file(channel: grpc.Channel, file_name: str, tags: List[str] = None, app_name: str = APP_NAME_FILE_SCAN) -> str:
+def scan_file(channel: grpc.Channel, file_name: str, tags: List[str] = None,
+              pml: bool = False, feedback: bool = False) -> str:
     try:
         f = open(file_name, "rb")
         fid = os.path.basename(file_name)
@@ -165,10 +193,10 @@ def scan_file(channel: grpc.Channel, file_name: str, tags: List[str] = None, app
         logger.debug("Permission error: " + str(err))
         raise AMaasException(AMaasErrorCode.MSG_ID_ERR_FILE_NO_PERMISSION, file_name)
 
-    return _scan_data(channel, f, n, fid, tags, app_name)
+    return _scan_data(channel, f, n, fid, tags, pml, feedback)
 
 
 def scan_buffer(channel: grpc.Channel, bytes_buffer: bytes, uid: str, tags: List[str] = None,
-                app_name: str = APP_NAME_FILE_SCAN) -> str:
+                pml: bool = False, feedback: bool = False) -> str:
     f = io.BytesIO(bytes_buffer)
-    return _scan_data(channel, f, len(bytes_buffer), uid, tags, app_name)
+    return _scan_data(channel, f, len(bytes_buffer), uid, tags, pml, feedback)