fix(file): use operation for get blocks

Author: davidt99

intezer_sdk/_api.py

@@ -515,6 +515,11 @@ def get_code_reuse_by_code_block(self, sha256: str):
             'POST', f'/files/{sha256}/code-reuse-by-code-block'
         )
 
+        if response.status_code == HTTPStatus.NOT_FOUND:
+            raise errors.HashDoesNotExistError(response)
+        if response.status_code == HTTPStatus.CONFLICT:
+            raise ValueError('sha256 is not a code item')
+
         raise_for_status(response)
 
         return response.json()['result_url']

intezer_sdk/file.py

@@ -1,17 +1,16 @@
+import datetime
 from dataclasses import dataclass
-from http import HTTPStatus
-from time import sleep
 from typing import IO
-from typing import Optional
 from typing import List
+from typing import Optional
 from typing import Union
 
+from intezer_sdk import _operation
 from intezer_sdk import consts
-from intezer_sdk import errors
+from intezer_sdk import operation
 from intezer_sdk._api import IntezerApi
-from intezer_sdk._api import raise_for_status
-from intezer_sdk.api import get_global_api
 from intezer_sdk.api import IntezerApiClient
+from intezer_sdk.api import get_global_api
 from intezer_sdk.index import Index
 
 
@@ -25,6 +24,7 @@ class Block:
     def is_common(self):
         return self.software_type == 'common'
 
+
 class File:
     """
     File is a class for file-related operations including indexing and downloading.
@@ -49,6 +49,7 @@ def __init__(self,
         self._sha256 = sha256
         self._api = IntezerApi(api or get_global_api())
         self._index: Optional[Index] = None
+        self._operations = {}
 
     @property
     def sha256(self) -> str:
@@ -154,40 +155,25 @@ def download(self,
 
         self._api.download_file_by_sha256(self._sha256, path, output_stream, password_protection)
 
-    def _get_result_from_task(self, result_url: str, sleep_time: int):
-        response = self._api.api.request_with_refresh_expired_access_token(
-            'GET', result_url)
-
-        if response.status_code == HTTPStatus.NOT_FOUND:
-            raise errors.HashDoesNotExistError(response)
-        if response.status_code == HTTPStatus.CONFLICT:
-            raise ValueError('sha256 is not a code item')
-
-        while response.status_code == HTTPStatus.ACCEPTED:
-            sleep(sleep_time)
-            response = self._api.api.request_with_refresh_expired_access_token(
-                'GET', result_url)
-        raise_for_status(response)
-        return response.json()['result']
-
-    def get_code_blocks(self, interval: int = consts.CHECK_STATUS_INTERVAL) -> List[Block]:
-        '''
+    def get_code_blocks(self,
+                        wait: Union[bool, int] = False,
+                        wait_timeout: Optional[datetime.timedelta] = None) -> operation.Operation:
+        """
         Retrieves a report containing information about reused code blocks for the given SHA-256 hash.
 
-        :param interval: The interval to wait between checks.
+        :param wait: Should wait until the operation completes.
+        :param wait_timeout: Maximum duration to wait for operation completion.
 
         Returns:
-            List[Block]: A sorted list of Block objects representing the code blocks found in the analysis.
-        '''
+            operation.Operation: An operation object that will contain the code blocks result.
+        """
         if not self._sha256:
             raise ValueError('Code block report is only supported for sha256-based files')
-        
+
         result_url = self._api.get_code_reuse_by_code_block(self._sha256)
-        # This endpoint acts different. We don't get a status and instead have to use
-        # the HTTP status code to wait for the report.
-        result = self._get_result_from_task(result_url, interval)
-        blocks: list[Block] = []
-        for address, block in result['blocks'].items():
-            blocks.append(
-                Block(int(address), block['software_type'], block['code_reuse']))
-        return sorted(blocks, key=lambda b: b.address)
+        return _operation.handle_operation(self._operations,
+                                           self._api,
+                                           'Code blocks',
+                                           result_url,
+                                           wait,
+                                           wait_timeout)

tests/unit/test_file.py

@@ -1,8 +1,6 @@
 import datetime
 import io
 from http import HTTPStatus
-import json
-import os
 from unittest.mock import mock_open
 from unittest.mock import patch
 
@@ -281,29 +279,104 @@ def test_download_to_path_succeeds(self):
                 mock_file.assert_called_once()
                 mock_file().write.assert_called_with(b'file_content')
 
-    def test_code_reuse_by_block(self):
-        TEST_HASH = '73c677dd3b264e7eb80e26e78ac9df1dba30915b5ce3b1bc1c83db52b9c6b30e'
-        
-        def load_response_json(file_name: str) -> dict:
-            path_to_file = os.path.join(os.path.dirname(__file__), '..', 'resources', file_name)
-            with open(path_to_file, 'rb') as file:
-                return json.load(file)
-    
+    def test_get_code_blocks_returns_operation_when_wait_is_true(self):
+        # Arrange
+        test_hash = '73c677dd3b264e7eb80e26e78ac9df1dba30915b5ce3b1bc1c83db52b9c6b30e'
+        result_url = f'/analyses/51ea282b-0542-4578-a44a-e60fdfb0d3ec/code-reuse-by-code-block'
+        result = {
+            'blocks': {
+                '101220': {
+                    'code_reuse': ['Common'],
+                    'software_type': 'common'
+                },
+                '101244': {
+                    'code_reuse': ['WannaCry'],
+                    'software_type': 'malware'
+                }
+            }
+        }
+
         with responses.RequestsMock() as mock:
             mock.post(
-                     url=consts.ANALYZE_URL +
-                     f'/api/v2-0/files/{TEST_HASH}/code-reuse-by-code-block',
-                     status=HTTPStatus.OK,
-                     json=load_response_json('code_reuse_block_response.json'))
+                url=f'{self.full_url}/files/{test_hash}/code-reuse-by-code-block',
+                status=HTTPStatus.OK,
+                json={'result_url': result_url})
             mock.get(
-                     url=consts.ANALYZE_URL +
-                     '/api/v2-0/analyses/51ea282b-0542-4578-a44a-e60fdfb0d3ec/code-reuse-by-code-block',
-                     status=HTTPStatus.OK,
-                     json=load_response_json('code_reuse_block_report.json'))
+                url=f'{self.full_url}{result_url}',
+                status=HTTPStatus.OK,
+                json={'status': 'succeeded', 'result': result})
+
+            file_object = File(sha256=test_hash)
+
+            # Act
+            operation = file_object.get_code_blocks(wait=True)
+
+        # Assert
+        self.assertEqual(operation.status, consts.AnalysisStatusCode.FINISHED)
+        self.assertIsNotNone(operation.result)
+        self.assertIn('blocks', operation.result)
+        self.assertEqual(len(operation.result['blocks']), 2)
+        self.assertEqual(operation.result['blocks']['101220']['software_type'], 'common')
+        self.assertEqual(operation.result['blocks']['101244']['software_type'], 'malware')
+
+    def test_get_code_blocks_for_file_path_raises_value_error(self):
+        # Arrange
+        file_obj = File(file_path='/path/to/file')
+
+        # Act + Assert
+        with self.assertRaises(ValueError):
+            file_obj.get_code_blocks()
+
+    def test_get_code_blocks_without_wait_returns_operation_in_progress(self):
+        # Arrange
+        test_hash = 'a' * 64
+        result_url = f'/analyses/test-id/code-reuse-by-code-block'
+
+        with responses.RequestsMock() as mock:
+            mock.post(
+                url=f'{self.full_url}/files/{test_hash}/code-reuse-by-code-block',
+                status=HTTPStatus.OK,
+                json={'result_url': result_url})
 
-            file_object = File(sha256=TEST_HASH)
-            report = file_object.get_code_blocks()
+            file_object = File(sha256=test_hash)
 
-            self.assertEqual(len(report), 2527)
-            self.assertEqual(len([x for x in report if x.is_common]), 1371)
-            self.assertEqual(len([x for x in report if x.software_type == 'malware']), 171)
+            # Act
+            operation = file_object.get_code_blocks()
+
+        # Assert
+        self.assertEqual(operation.status, consts.AnalysisStatusCode.IN_PROGRESS)
+        self.assertIsNone(operation.result)
+
+    def test_get_code_blocks_raises_hash_does_not_exist_error_when_file_not_found(self):
+        # Arrange
+        test_hash = 'a' * 64
+
+        with responses.RequestsMock() as mock:
+            mock.post(
+                url=f'{self.full_url}/files/{test_hash}/code-reuse-by-code-block',
+                status=HTTPStatus.NOT_FOUND,
+                json={'error': 'File not found'})
+
+            file_object = File(sha256=test_hash)
+
+            # Act + Assert
+            with self.assertRaises(errors.HashDoesNotExistError):
+                file_object.get_code_blocks()
+
+    def test_get_code_blocks_raises_value_error_when_file_is_not_code_item(self):
+        # Arrange
+        test_hash = 'a' * 64
+
+        with responses.RequestsMock() as mock:
+            mock.post(
+                url=f'{self.full_url}/files/{test_hash}/code-reuse-by-code-block',
+                status=HTTPStatus.CONFLICT,
+                json={'error': 'Not a code item'})
+
+            file_object = File(sha256=test_hash)
+
+            # Act + Assert
+            with self.assertRaises(ValueError) as context:
+                file_object.get_code_blocks()
+            
+            self.assertEqual(str(context.exception), 'sha256 is not a code item')