diff --git a/src/urlscan/__init__.py b/src/urlscan/__init__.py index 0dacbaa..1e32059 100644 --- a/src/urlscan/__init__.py +++ b/src/urlscan/__init__.py @@ -8,3 +8,5 @@ from .client import Client # noqa: F401 from .error import APIError, RateLimitError # noqa: F401 from .iterator import SearchIterator # noqa: F401 +from .pro import Pro # noqa: F401 +from .types import LiveScanResourceType, VisibilityType # noqa: F401 diff --git a/src/urlscan/client.py b/src/urlscan/client.py index 8dbd2a8..4ad2ee3 100644 --- a/src/urlscan/client.py +++ b/src/urlscan/client.py @@ -95,7 +95,7 @@ class RateLimitMemo(TypedDict): RateLimitKey = Literal["public", "private", "unlisted", "retrieve", "search"] -class Client: +class BaseClient: def __init__( self, api_key: str, @@ -268,6 +268,44 @@ def post( req = session.build_request("POST", path, json=json, data=data) return self._send_request(session, req) + def put( + self, + path: str, + json: Any | None = None, + data: RequestData | None = None, + ) -> ClientResponse: + """Send a PUT request to a given API endpoint. + + Args: + path (str): Path. + json (Any | None, optional): Dict to send in request body as JSON. Defaults to None. + data (RequestData | None, optional): Dict to send in request body. Defaults to None. + + Returns: + ClientResponse: Response. + """ + session = self._get_session() + req = session.build_request("PUT", path, json=json, data=data) + return self._send_request(session, req) + + def delete( + self, + path: str, + params: QueryParamTypes | None = None, + ) -> ClientResponse: + """Send a DELETE request to a given API endpoint. + + Args: + path (str): Path. + params (QueryParamTypes | None, optional): Query parameters. Defaults to None. + + Returns: + ClientResponse: Response. + """ + session = self._get_session() + req = session.build_request("DELETE", path, params=params) + return self._send_request(session, req) + def download( self, path: str, @@ -296,6 +334,54 @@ def get_text(self, path: str, params: QueryParamTypes | None = None) -> str: res = self.get(path, params=params) return self._response_to_str(res) + def _get_error(self, res: ClientResponse) -> APIError | None: + try: + res.raise_for_status() + except httpx.HTTPStatusError as exc: + data: dict = exc.response.json() + message: str = data["message"] + description: str | None = data.get("description") + status: int = data["status"] + + # ref. https://urlscan.io/docs/api/#ratelimit + if status == 429: + rate_limit_reset_after = float( + exc.response.headers.get("X-Rate-Limit-Reset-After", 0) + ) + return RateLimitError( + message, + description=description, + status=status, + rate_limit_reset_after=rate_limit_reset_after, + ) + + return APIError(message, description=description, status=status) + + return None + + def _response_to_json(self, res: ClientResponse) -> dict: + error = self._get_error(res) + if error: + raise error + + return res.json() + + def _response_to_str(self, res: ClientResponse) -> str: + error = self._get_error(res) + if error: + raise error + + return res.text + + def _response_to_content(self, res: ClientResponse) -> bytes: + error = self._get_error(res) + if error: + raise error + + return res.content + + +class Client(BaseClient): def get_result(self, uuid: str) -> dict: """Get a result of a scan by UUID. @@ -607,49 +693,3 @@ def mapping(res_or_error: dict | Exception) -> dict | Exception: return self.get_result(uuid) return [(url, mapping(res_or_error)) for url, res_or_error in responses] - - def _get_error(self, res: ClientResponse) -> APIError | None: - try: - res.raise_for_status() - except httpx.HTTPStatusError as exc: - data: dict = exc.response.json() - message: str = data["message"] - description: str | None = data.get("description") - status: int = data["status"] - - # ref. https://urlscan.io/docs/api/#ratelimit - if status == 429: - rate_limit_reset_after = float( - exc.response.headers.get("X-Rate-Limit-Reset-After", 0) - ) - return RateLimitError( - message, - description=description, - status=status, - rate_limit_reset_after=rate_limit_reset_after, - ) - - return APIError(message, description=description, status=status) - - return None - - def _response_to_json(self, res: ClientResponse) -> dict: - error = self._get_error(res) - if error: - raise error - - return res.json() - - def _response_to_str(self, res: ClientResponse) -> str: - error = self._get_error(res) - if error: - raise error - - return res.text - - def _response_to_content(self, res: ClientResponse) -> bytes: - error = self._get_error(res) - if error: - raise error - - return res.content diff --git a/src/urlscan/pro/__init__.py b/src/urlscan/pro/__init__.py new file mode 100644 index 0000000..0dcee06 --- /dev/null +++ b/src/urlscan/pro/__init__.py @@ -0,0 +1,49 @@ +from urlscan.client import BASE_URL, USER_AGENT, BaseClient, TimeoutTypes + +from .livescan import LiveScan + + +class Pro(BaseClient): + def __init__( + self, + api_key: str, + base_url: str = BASE_URL, + user_agent: str = USER_AGENT, + trust_env: bool = False, + timeout: TimeoutTypes = 60, + proxy: str | None = None, + verify: bool = True, + retry: bool = False, + ): + """ + Args: + api_key (str): Your urlscan.io API key. + base_url (str, optional): Base URL. Defaults to BASE_URL. + user_agent (str, optional): User agent. Defaults to USER_AGENT. + trust_env (bool, optional): Enable or disable usage of environment variables for configuration. Defaults to False. + timeout (TimeoutTypes, optional): timeout configuration to use when sending request. Defaults to 60. + proxy (str | None, optional): Proxy URL where all the traffic should be routed. Defaults to None. + verify (bool, optional): Either `True` to use an SSL context with the default CA bundle, `False` to disable verification. Defaults to True. + retry (bool, optional): Whether to use automatic X-Rate-Limit-Reset-After HTTP header based retry. Defaults to False. + """ + super().__init__( + api_key, + base_url=base_url, + user_agent=user_agent, + trust_env=trust_env, + timeout=timeout, + proxy=proxy, + verify=verify, + retry=retry, + ) + + self.livescan = LiveScan( + api_key=api_key, + base_url=base_url, + user_agent=user_agent, + trust_env=trust_env, + timeout=timeout, + proxy=proxy, + verify=verify, + retry=retry, + ) diff --git a/src/urlscan/pro/livescan.py b/src/urlscan/pro/livescan.py new file mode 100644 index 0000000..af69a68 --- /dev/null +++ b/src/urlscan/pro/livescan.py @@ -0,0 +1,191 @@ +from typing import Any + +from urlscan.client import BaseClient, _compact +from urlscan.types import LiveScanResourceType, VisibilityType + + +class LiveScan(BaseClient): + def get_scanners(self) -> dict: + """Get a list of available Live Scanning nodes along with their current metadata. + + Returns: + dict: List of available scanners with metadata. + + Reference: + https://docs.urlscan.io/apis/urlscan-openapi/live-scanning/livescanscanners + """ + return self.get_json("/api/v1/livescan/scanners/") + + def task( + self, + url: str, + *, + scanner_id: str, + visibility: VisibilityType | None = None, + page_timeout: int | None = None, + capture_delay: int | None = None, + extra_headers: dict[str, str] | None = None, + enable_features: list[str] | None = None, + disable_features: list[str] | None = None, + ) -> dict: + """Task a URL to be scanned. + + The HTTP request will return with the scan UUID immediately and then it is your responsibility to poll the result resource type until the scan has finished. + + Args: + url (str): URL to scan. + scanner_id (str): Scanner ID (e.g., "de01" for Germany). + visibility (VisibilityType | None, optional): Visibility of the scan. Defaults to None. + page_timeout (int | None, optional): Time to wait for the whole scan process (in ms). Defaults to None. + capture_delay (int | None, optional): Delay after page load before capturing (in ms). Defaults to None. + extra_headers (dict[str, str] | None, optional): Extra HTTP headers. Defaults to None. + enable_features (list[str] | None, optional): Features to enable. Defaults to None. + disable_features (list[str] | None, optional): Features to disable. Defaults to None. + + Returns: + dict: Response containing the scan UUID. + + Reference: + https://docs.urlscan.io/apis/urlscan-openapi/live-scanning/livescantask + """ + task: dict[str, Any] = _compact( + { + "url": url, + "visibility": visibility, + } + ) + scanner: dict[str, Any] = _compact( + { + "pageTimeout": page_timeout, + "captureDelay": capture_delay, + "extraHeaders": extra_headers, + "enableFeatures": enable_features, + "disableFeatures": disable_features, + } + ) + data: dict[str, Any] = _compact({"task": task, "scanner": scanner}) + + res = self.post(f"/api/v1/livescan/{scanner_id}/task/", json=data) + return self._response_to_json(res) + + def scan( + self, + url: str, + *, + scanner_id: str, + visibility: VisibilityType | None = None, + page_timeout: int | None = None, + capture_delay: int | None = None, + extra_headers: dict[str, str] | None = None, + enable_features: list[str] | None = None, + disable_features: list[str] | None = None, + ) -> dict: + """Task a URL to be scanned. The HTTP request will block until the scan has finished. + + Args: + url (str): URL to scan. + scanner_id (str): Scanner ID (e.g., "de01" for Germany). + visibility (VisibilityType | None, optional): Visibility of the scan. Defaults to None. + page_timeout (int | None, optional): Time to wait for the whole scan process (in ms). Defaults to None. + capture_delay (int | None, optional): Delay after page load before capturing (in ms). Defaults to None. + extra_headers (dict[str, str] | None, optional): Extra HTTP headers. Defaults to None. + enable_features (list[str] | None, optional): Features to enable. Defaults to None. + disable_features (list[str] | None, optional): Features to disable. Defaults to None. + + Returns: + dict: Response containing the scan UUID. + + Reference: + https://docs.urlscan.io/apis/urlscan-openapi/live-scanning/livescanscan + """ + task: dict[str, Any] = _compact( + { + "url": url, + "visibility": visibility, + } + ) + scanner: dict[str, Any] = _compact( + { + "pageTimeout": page_timeout, + "captureDelay": capture_delay, + "extraHeaders": extra_headers, + "enableFeatures": enable_features, + "disableFeatures": disable_features, + } + ) + data: dict[str, Any] = _compact({"task": task, "scanner": scanner}) + + res = self.post(f"/api/v1/livescan/{scanner_id}/scan/", json=data) + return self._response_to_json(res) + + def get_resource( + self, + *, + scanner_id: str, + resource_type: LiveScanResourceType, + resource_id: str, + ) -> Any: + """Get the historical observations for a specific hostname in the "Hostnames" data source. + + Args: + scanner_id (str): Scanner ID (e.g., "de01" for Germany). + resource_type (LiveScanResourceType): Type of resource ("result", "screenshot", "dom", "response", or "download"). + resource_id (str): Resource ID. For result/screenshot/dom: UUID of the scan. For response/download: SHA256 of the resource. + + Returns: + Any: Resource content. Returns dict for "result", str for "dom", bytes for binary resources. + + Reference: + https://docs.urlscan.io/apis/urlscan-openapi/live-scanning/livescangetresource + """ + path = f"/api/v1/livescan/{scanner_id}/{resource_type}/{resource_id}" + + if resource_type == "result": + return self.get_json(path) + + if resource_type in ("screenshot", "response", "download"): + return self.get_content(path) + + if resource_type == "dom": + return self.get_text(path) + + return self.get(path) + + def store( + self, + *, + scanner_id: str, + scan_id: str, + visibility: VisibilityType, + ) -> dict: + """Store the temporary scan as a permanent snapshot on urlscan.io. + + Args: + scanner_id (str): Scanner ID (e.g., "de01" for Germany). + scan_id (str): Scan UUID. + visibility (VisibilityType): Visibility for the stored scan ("public", "private", or "unlisted"). + + Reference: + https://docs.urlscan.io/apis/urlscan-openapi/live-scanning/livescanstore + """ + data = {"task": {"visibility": visibility}} + res = self.put(f"/api/v1/livescan/{scanner_id}/{scan_id}/", json=data) + return self._response_to_json(res) + + def purge( + self, + *, + scanner_id: str, + scan_id: str, + ) -> dict: + """Purge temporary scan from scanner immediately. Scans will be automatically purged after 60 minutes. + + Args: + scanner_id (str): Scanner ID (e.g., "de01" for Germany). + scan_id (str): Scan UUID. + + Reference: + https://docs.urlscan.io/apis/urlscan-openapi/live-scanning/livescandiscard + """ + res = self.delete(f"/api/v1/livescan/{scanner_id}/{scan_id}/") + return self._response_to_json(res) diff --git a/src/urlscan/types.py b/src/urlscan/types.py index 2074645..b828698 100644 --- a/src/urlscan/types.py +++ b/src/urlscan/types.py @@ -4,3 +4,4 @@ SearchType = Literal["search"] RetrieveType = Literal["retrieve"] ActionType = VisibilityType | SearchType | RetrieveType +LiveScanResourceType = Literal["result", "screenshot", "dom", "response", "download"] diff --git a/tests/conftest.py b/tests/conftest.py index e69de29..ecd3333 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -0,0 +1,25 @@ +import pytest +from pytest_httpserver import HTTPServer + +from urlscan import Client, Pro + + +@pytest.fixture +def api_key(): + return "dummy" + + +@pytest.fixture +def client(httpserver: HTTPServer, api_key: str): + with Client( + api_key=api_key, base_url=f"http://{httpserver.host}:{httpserver.port}" + ) as client: + yield client + + +@pytest.fixture +def pro(httpserver: HTTPServer, api_key: str): + with Pro( + api_key=api_key, base_url=f"http://{httpserver.host}:{httpserver.port}" + ) as client: + yield client diff --git a/tests/pro/__init__.py b/tests/pro/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/pro/test_livescan.py b/tests/pro/test_livescan.py new file mode 100644 index 0000000..4a18a94 --- /dev/null +++ b/tests/pro/test_livescan.py @@ -0,0 +1,46 @@ +from typing import Any + +from pytest_httpserver import HTTPServer + +from urlscan.pro import Pro + + +def test_get_scanners(pro: Pro, httpserver: HTTPServer): + data: dict[str, Any] = {"scanners": []} + httpserver.expect_request("/api/v1/livescan/scanners/").respond_with_json(data) + + got = pro.livescan.get_scanners() + assert got == data + + +def test_task(pro: Pro, httpserver: HTTPServer): + data = {"uuid": "dummy-uuid"} + httpserver.expect_request( + "/api/v1/livescan/de01/task/", + method="POST", + ).respond_with_json(data) + + got = pro.livescan.task(scanner_id="de01", url="http://example.com") + assert got == data + + +def test_scan(pro: Pro, httpserver: HTTPServer): + data = {"uuid": "dummy-uuid"} + httpserver.expect_request( + "/api/v1/livescan/de01/scan/", + method="POST", + ).respond_with_json(data) + + got = pro.livescan.scan(scanner_id="de01", url="http://example.com") + assert got == data + + +def test_purge(pro: Pro, httpserver: HTTPServer): + data = {"status": "purged"} + httpserver.expect_request( + "/api/v1/livescan/de01/dummy-uuid/", + method="DELETE", + ).respond_with_json(data) + + got = pro.livescan.purge(scanner_id="de01", scan_id="dummy-uuid") + assert got == data diff --git a/tests/test_client.py b/tests/test_client.py index 5bcd63e..f42f491 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -11,19 +11,6 @@ from urlscan.error import RateLimitError, RateLimitRemainingError -@pytest.fixture -def api_key(): - return "dummy" - - -@pytest.fixture -def client(httpserver: HTTPServer, api_key: str): - with Client( - api_key=api_key, base_url=f"http://{httpserver.host}:{httpserver.port}" - ) as client: - yield client - - def test_get(client: Client, httpserver: HTTPServer): data = {"foo": "bar"} httpserver.expect_request("/dummy").respond_with_json(data)