Poll cached full scan results until ready

socketdev/core/api.py

@@ -1,4 +1,5 @@
 import base64
+from typing import Optional
 from socketdev.log import log
 
 import requests
@@ -47,6 +48,8 @@ def do_request(
         payload: [dict, str] = None,
         files: list = None,
         method: str = "GET",
+        *,
+        request_timeout: Optional[float] = None,
     ) -> Response:
         if self.encoded_key is None or self.encoded_key == "":
             raise APIKeyMissing
@@ -68,7 +71,8 @@ def format_headers(headers_dict):
 
             response = requests.request(
                 method.upper(), url, headers=headers, data=payload, files=files, 
-                timeout=self.request_timeout, verify=not self.allow_unverified
+                timeout=self.request_timeout if request_timeout is None else request_timeout,
+                verify=not self.allow_unverified
             )
             request_duration = time.time() - start_time
 

socketdev/fullscans/__init__.py

@@ -4,9 +4,12 @@
 from typing import Any, Dict, List, Optional, Union
 from dataclasses import dataclass, asdict, field
 import urllib.parse
+import math
+import time
 from ..core.dedupe import Dedupe
 from ..core.enums import unknown_enum_value
 from ..utils import IntegrationType, Utils
+from ..exceptions import APIFailure, APITimeout
 
 log = logging.getLogger("socketdev")
 
@@ -960,15 +963,71 @@ def stream_diff(
             )
         return {}
 
-    def stream(self, org_slug: str, full_scan_id: str, use_types: bool = False) -> Union[dict, FullScanStreamResponse]:
-        path = "orgs/" + org_slug + "/full-scans/" + full_scan_id
-        response = self.api.do_request(path=path, method="GET")
+    def stream(
+        self,
+        org_slug: str,
+        full_scan_id: str,
+        use_types: bool = False,
+        *,
+        poll_timeout_seconds: Optional[float] = None,
+        poll_initial_interval_seconds: float = 5.0,
+        poll_max_interval_seconds: float = 10.0,
+    ) -> Union[dict, FullScanStreamResponse]:
+        request_timeout = self.api.request_timeout
+        timeout = request_timeout if poll_timeout_seconds is None else poll_timeout_seconds
+        values = (timeout, poll_initial_interval_seconds, poll_max_interval_seconds, request_timeout)
+        if any(
+            isinstance(value, bool)
+            or not isinstance(value, (int, float))
+            or not math.isfinite(value)
+            or value <= 0
+            for value in values
+        ):
+            raise ValueError("Full scan polling timeout and intervals must be positive finite numbers")
+        if poll_initial_interval_seconds > poll_max_interval_seconds:
+            raise ValueError("Initial polling interval must not exceed maximum polling interval")
+
+        path = (
+            "orgs/" + urllib.parse.quote(org_slug, safe="") + "/full-scans/"
+            + urllib.parse.quote(full_scan_id, safe="") + "?cached=true"
+        )
+        deadline = time.monotonic() + timeout
+        interval = poll_initial_interval_seconds
+        while True:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise APITimeout("Timed out waiting for full scan results")
+            try:
+                response = self.api.do_request(
+                    path=path,
+                    method="GET",
+                    request_timeout=min(30.0, float(request_timeout), remaining),
+                )
+            except APIFailure as error:
+                if not error.is_transient_error():
+                    raise
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise APITimeout("Timed out waiting for full scan results") from error
+                time.sleep(min(interval, remaining))
+                interval = min(poll_max_interval_seconds, interval * 1.5)
+                continue
+            if response.status_code == 202:
+                response.close()
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise APITimeout("Timed out waiting for full scan results")
+                time.sleep(min(interval, remaining))
+                interval = min(poll_max_interval_seconds, interval * 1.5)
+                continue
+            break
 
         if response.status_code == 200:
             try:
                 stream_str = []
                 artifacts = {}
                 result = response.text
+                response.close()
                 result = result.strip('"').strip()
                 for line in result.split("\n"):
                     if line != '"' and line != "" and line is not None:
@@ -993,6 +1052,7 @@ def stream(self, org_slug: str, full_scan_id: str, use_types: bool = False) -> U
                 return artifacts
 
             except Exception as e:
+                response.close()
                 error_message = f"Error parsing stream response: {str(e)}"
                 log.error(error_message)
                 if use_types:
@@ -1001,7 +1061,10 @@ def stream(self, org_slug: str, full_scan_id: str, use_types: bool = False) -> U
                     )
                 return {}
 
-        error_message = response.json().get("error", {}).get("message", "Unknown error")
+        try:
+            error_message = response.json().get("error", {}).get("message", "Unknown error")
+        finally:
+            response.close()
         log.error(f"Error streaming full scan: {response.status_code}, message: {error_message}")
         if use_types:
             return FullScanStreamResponse.from_dict(

socketdev/version.py

@@ -1 +1 @@
-__version__ = "3.6.0"
+__version__ = "3.7.0"

tests/support/socket_api_fixture.py

@@ -0,0 +1,121 @@
+"""Reusable localhost HTTP fixture for real Socket SDK transport tests."""
+
+from __future__ import annotations
+
+import json
+import threading
+from dataclasses import dataclass
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from typing import Callable, Dict, List, Optional, Tuple
+from urllib.parse import parse_qs, urlsplit
+
+
+@dataclass
+class SocketRequest:
+    method: str
+    path: str
+    query: Dict[str, List[str]]
+    headers: Dict[str, str]
+    body: bytes
+
+
+@dataclass
+class SocketResponse:
+    status: int
+    payload: object
+    raw_body: Optional[bytes] = None
+
+
+Route = Callable[[SocketRequest], SocketResponse]
+
+
+class SocketApiFixture:
+    def __init__(self) -> None:
+        self.requests: List[SocketRequest] = []
+        self._routes: Dict[Tuple[str, str], Route] = {}
+        self._server: Optional[ThreadingHTTPServer] = None
+        self._thread: Optional[threading.Thread] = None
+
+    @property
+    def base_url(self) -> str:
+        if self._server is None:
+            raise RuntimeError("Socket API fixture is not running")
+        host, port = self._server.server_address
+        return f"http://{host}:{port}/v0"
+
+    def route(
+        self,
+        method: str,
+        path: str,
+        response: SocketResponse | Route,
+    ) -> None:
+        if callable(response):
+            handler = response
+        else:
+
+            def handler(_request: SocketRequest) -> SocketResponse:
+                return response
+
+        self._routes[(method.upper(), path)] = handler
+
+    def start(self) -> None:
+        fixture = self
+
+        class Handler(BaseHTTPRequestHandler):
+            def do_GET(self) -> None:
+                self._dispatch()
+
+            def do_POST(self) -> None:
+                self._dispatch()
+
+            def _dispatch(self) -> None:
+                parsed = urlsplit(self.path)
+                length = int(self.headers.get("Content-Length", "0"))
+                request = SocketRequest(
+                    method=self.command,
+                    path=parsed.path.removeprefix("/v0/"),
+                    query=parse_qs(parsed.query, keep_blank_values=True),
+                    headers=dict(self.headers.items()),
+                    body=self.rfile.read(length),
+                )
+                fixture.requests.append(request)
+                route = fixture._routes.get((request.method, request.path))
+                if route is None and request.path == "organizations":
+                    response = SocketResponse(
+                        200,
+                        {"organizations": {"org-fixture": {"slug": "fixture-org"}}},
+                    )
+                elif route is None:
+                    response = SocketResponse(404, {"error": {"message": "Not found"}})
+                else:
+                    response = route(request)
+                body = response.raw_body
+                if body is None:
+                    body = json.dumps(response.payload).encode("utf-8")
+                self.send_response(response.status)
+                self.send_header("Content-Type", "application/json")
+                self.send_header("Content-Length", str(len(body)))
+                self.end_headers()
+                self.wfile.write(body)
+
+            def log_message(self, _format: str, *args: object) -> None:
+                return
+
+        self._server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
+        self._thread.start()
+
+    def close(self) -> None:
+        if self._server is not None:
+            self._server.shutdown()
+            self._server.server_close()
+            self._server = None
+        if self._thread is not None:
+            self._thread.join(timeout=2)
+            self._thread = None
+
+
+def socket_api_fixture() -> SocketApiFixture:
+    fixture = SocketApiFixture()
+    fixture.start()
+    return fixture