Add fetching content from content_url

agent/semgrep_agent.py

@@ -79,7 +79,7 @@ def process(self, message: m.Message) -> None:
             message: A message containing the path and the content of the file to be processed
 
         """
-        content = message.data.get("content")
+        content = utils.get_file_content(message)
         path = message.data.get("path")
 
         if content is None:

agent/utils.py

@@ -1,16 +1,24 @@
 """Utilities for Semgrep Agent"""
 import dataclasses
 import mimetypes
+import requests
+import logging
 import os
 import re
 from typing import Any, Iterator
 from urllib import parse
 
+
+import tenacity
 import magic
 from ostorlab.agent.kb import kb
 from ostorlab.agent.mixins import agent_report_vulnerability_mixin
+from ostorlab.agent.message import message as m
+
 
 LINE_SIZE_MAX = 5000
+DOWNLOAD_REQUEST_TIMEOUT = 60
+NUMBER_RETRIES = 3
 
 RISK_RATING_MAPPING = {
     "UNKNOWN": agent_report_vulnerability_mixin.RiskRating.POTENTIALLY,
@@ -19,6 +27,8 @@
     "HIGH": agent_report_vulnerability_mixin.RiskRating.HIGH,
 }
 
+logger = logging.getLogger(__name__)
+
 
 @dataclasses.dataclass
 class Vulnerability:
@@ -138,3 +148,49 @@ def get_file_type(content: bytes, path: str | None) -> str:
         if len(file_split) < 2:
             return get_file_type(content, None)
         return file_split
+
+
+@tenacity.retry(
+    stop=tenacity.stop_after_attempt(NUMBER_RETRIES),
+    retry=tenacity.retry_if_exception_type(requests.exceptions.RequestException),
+    retry_error_callback=lambda retry_state: retry_state.outcome.result()
+    if retry_state.outcome is not None
+    else None,
+)
+def _download_file(file_url: str) -> bytes | None:
+    """Download a file.
+
+    Args:
+        file_url : The URL of the file to download.
+
+    Returns:
+        bytes: The content of the file.
+    """
+    response = requests.get(file_url, timeout=DOWNLOAD_REQUEST_TIMEOUT)
+    if response.status_code == 200 and isinstance(response.content, bytes) is True:
+        return response.content
+
+    return None
+
+
+def get_file_content(message: m.Message) -> bytes | None:
+    """Get the file content from a message.
+
+    Args:
+        message : The message containing the file data.
+
+    Returns:
+        bytes: The content of the file.
+    """
+    content = message.data.get("content")
+    if isinstance(content, bytes) is True:
+        return content
+    content_url: str | None = message.data.get("content_url")
+    if content_url is not None:
+        try:
+            content = _download_file(content_url)
+        except requests.exceptions.RequestException as e:
+            logger.error("Could not download file %s. Error: %s.", content_url, e)
+        return content
+
+    return None

requirement.txt

@@ -2,4 +2,5 @@ ostorlab[agent]
 rich
 semgrep
 python-magic
-jsbeautifier
+jsbeautifier
+tenacity

tests/test-requirement.txt

@@ -1,5 +1,7 @@
 pytest
 pytest-mock
+requests_mock
+types-requests
 ruff
 mypy
 typing-extensions

tests/utils_test.py

@@ -1,8 +1,10 @@
 """Unittests for Semgrep Agent Utilities"""
 from typing import Any
+import requests
 
 import pytest
-from ostorlab.agent.message import message
+import requests_mock as reqs_mock
+from ostorlab.agent.message import message as m
 
 from agent import utils
 
@@ -81,7 +83,7 @@ def testParseResults_whenNoVulnerabilitiesAreFound_returnsVulnerability(
 
 
 def testGetFileType_withPathProvided_returnsFileType(
-    scan_message_file: message.Message,
+    scan_message_file: m.Message,
 ) -> None:
     """Unittest for the file type extraction:
     case when the path is provided
@@ -94,7 +96,7 @@ def testGetFileType_withPathProvided_returnsFileType(
 
 
 def testGetFileType_withoutPathProvided_returnsFileType(
-    scan_message_file: message.Message,
+    scan_message_file: m.Message,
 ) -> None:
     """Unittest for the file type extraction:
     case when the path is not provided
@@ -151,3 +153,58 @@ def testFilterDescription_caseRegexRedos_returnFilteredDescription() -> None:
         "library such as https://www.npmjs.com/package/recheck to verify that the regex does not appear "
         "vulnerable to ReDoS."
     )
+
+
+def testGetFileContent_whenContentIsNotNone_returnTheContent() -> None:
+    """Test that the content is returned when it is not empty."""
+    message = m.Message.from_data(
+        selector="v3.asset.file.android.apk", data={"content": b"content"}
+    )
+
+    content = utils.get_file_content(message)
+
+    assert content == b"content"
+
+
+def testGetFileContent_whenContentIsEmpty_shouldTryToDownloadTheFile(
+    requests_mock: reqs_mock.mocker.Mocker,
+) -> None:
+    """Test that the file is downloaded when the content is empty."""
+    message = m.Message.from_data(
+        selector="v3.asset.file.android.apk",
+        data={"content_url": "https://example.com/legendary.apk"},
+    )
+    requests_mock.get(
+        "https://example.com/legendary.apk", content=b"downloaded_content"
+    )
+
+    content = utils.get_file_content(message)
+
+    assert content == b"downloaded_content"
+
+
+def testGetFileContent_whenContentUrlIsUnreacheable_shouldRetryThreeTimes(
+    requests_mock: reqs_mock.mocker.Mocker,
+) -> None:
+    """Test that the file download is retried three times."""
+    message = m.Message.from_data(
+        selector="v3.asset.file.android.apk",
+        data={"content_url": "https://example.com/legendary.apk"},
+    )
+    download_file_mock = requests_mock.get(
+        "https://example.com/legendary.apk", exc=requests.exceptions.ConnectionError
+    )
+
+    content = utils.get_file_content(message)
+
+    assert content is None
+    assert download_file_mock.call_count == 3
+
+
+def testGetFileContent_whenNoContentIsAvailable_shouldReturnNone() -> None:
+    """Test that None is returned when no content is available."""
+    message = m.Message.from_data(selector="v3.asset.file.android.apk", data={})
+
+    content = utils.get_file_content(message)
+
+    assert content is None