From 95526eb2e4c2ccb569517294df9f64c3a553ccc8 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 00:33:33 +0000 Subject: [PATCH] fix: follow pagination in get_profiles, get_users, and get_user_apikeys These methods returned only the API's first page (50 items by default) and dropped the `next` offset, so callers had no way to get the rest. They now follow `next` in pages of up to 200 and return every item; `limit` caps the total and `offset` starts from an earlier page. New `profiles_page` and `users_page` methods return a single `base.ListPage` (items plus `next`), and `iter_profiles` and `iter_users` stream items across pages. FakeTriage now pages the way the docs describe (default 50, max 200), and pages API key lists too. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KBLBzm6b7WftoJvKs7VbJw --- CHANGELOG.md | 3 ++ src/pyhatching/__init__.py | 102 ++++++++++++++++++++++++++++++------- src/pyhatching/base.py | 11 ++++ src/pyhatching/testing.py | 5 +- tests/test_client.py | 58 +++++++++++++++++++++ 5 files changed, 159 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cb4befc..e85ffad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,9 @@ ## Unreleased ### Features +- Add `profiles_page`, `iter_profiles`, `users_page`, and `iter_users`, and a `base.ListPage` model for a single page of a paginated collection (with its `next` offset). ### Fixes +- `get_profiles`, `get_users`, and `get_user_apikeys` now follow the API's pagination and return every item (`limit` caps how many), instead of only the first page with no way to get the next one. +- `FakeTriage` pages like the API: a default page size of 50, at most 200, and API key lists are paginated too. - Match the live API in `pyhatching.testing`'s task report: DNS record types include their class (`"IN A"`), and `orig` is set on every process in the sample's tree. ## v1.0.0 - 2026-11-01 diff --git a/src/pyhatching/__init__.py b/src/pyhatching/__init__.py index 4a6e1eb..71c34af 100644 --- a/src/pyhatching/__init__.py +++ b/src/pyhatching/__init__.py @@ -659,10 +659,14 @@ async def search_page( params = {"query": query, "limit": min(limit, MAX_PAGE_SIZE), "offset": offset} return await self._model(base.SamplesPage, "get", "/v0/search", params=params) - async def _pages(self, getpage, limit: int) -> AsyncIterator[base.SamplesPage | base.ErrorResponse]: - """Yield pages from ``getpage(limit, offset)`` until ``limit`` items (0 is no limit).""" + async def _pages( + self, getpage, limit: int, offset: int | str | None = None + ) -> AsyncIterator[base.SamplesPage | base.ListPage | base.ErrorResponse]: + """Yield pages from ``getpage(limit, offset)`` until ``limit`` items (0 is no limit). + + Starts at ``offset`` (the ``next`` value of an earlier page), or the first page. + """ - offset = None count = 0 while True: size = MAX_PAGE_SIZE if not limit else min(MAX_PAGE_SIZE, limit - count) @@ -686,14 +690,32 @@ async def _iter(self, getpage, limit: int) -> AsyncIterator[base.SampleInfo]: if limit and count >= limit: return - async def _list(self, getpage, limit: int) -> list[base.SampleInfo] | base.ErrorResponse: + async def _list( + self, getpage, limit: int, offset: int | str | None = None + ) -> list | base.ErrorResponse: ret = [] - async for page in self._pages(getpage, limit): + async for page in self._pages(getpage, limit, offset): if isinstance(page, base.ErrorResponse): return page ret.extend(page.data) return ret[:limit] if limit else ret + async def _list_page( + self, model, uri: str, limit: int = MAX_PAGE_SIZE, offset: int | str | None = None + ) -> base.ListPage | base.ErrorResponse: + """Get a single page of a paginated collection, its items converted to ``model``.""" + + params = {"limit": min(limit, MAX_PAGE_SIZE), "offset": offset} + resp, resp_dict = await self._request("get", uri, params=params) + page = self.convert_resp(base.ListPage, resp, resp_dict) + if isinstance(page, base.ErrorResponse): + return page + page.data = [ + item if isinstance(item, base.TriageModel) + else self.convert_resp(model, resp, item) for item in page.data + ] + return page + async def iter_samples( self, subset: enums.SampleSubsets | str = enums.SampleSubsets.OWNED, limit: int = 20 ) -> AsyncIterator[base.SampleInfo]: @@ -1253,10 +1275,34 @@ async def get_profile( base.HatchingProfileResponse, "get", f"/v0/profiles/{_seg(profile_id)}" ) + async def profiles_page( + self, limit: int = MAX_PAGE_SIZE, offset: int | str | None = None + ) -> base.ListPage | base.ErrorResponse: + """Get a single page of your sandbox analysis profiles (GET /v0/profiles). + + ``data`` holds ``base.HatchingProfileResponse`` items. ``offset`` is the + ``next`` value of the previous page. + """ + + return await self._list_page(base.HatchingProfileResponse, "/v0/profiles", limit, offset) + + async def iter_profiles(self, limit: int = 0) -> AsyncIterator[base.HatchingProfileResponse]: + """Yield up to ``limit`` of your sandbox analysis profiles (0 for all of them).""" + + async for profile in self._iter(self.profiles_page, limit): + yield profile + async def get_profiles( self, limit: int | None = None, offset: int | str | None = None ) -> list[base.HatchingProfileResponse] | base.ErrorResponse: - """Get all sandbox analysis profiles for your account. + """Get your sandbox analysis profiles, following the API's pagination. + + Parameters + ---------- + limit : int | None, optional + The most profiles to return, None or 0 for all of them. + offset : int | str | None, optional + Start after a previous page (its ``next`` value), see ``profiles_page``. Returns ------- @@ -1266,12 +1312,7 @@ async def get_profiles( If there was an error. """ - return await self._model( - base.HatchingProfileResponse, - "get", - "/v0/profiles", - params={"limit": limit, "offset": offset}, - ) + return await self._list(self.profiles_page, limit or 0, offset) async def submit_profile( self, @@ -1448,14 +1489,34 @@ async def delete_rule(self, name: str) -> base.ErrorResponse | None: # ---- Users (private instances) ---- + async def users_page( + self, limit: int = MAX_PAGE_SIZE, offset: int | str | None = None + ) -> base.ListPage | base.ErrorResponse: + """Get a single page of your company's users (GET /v0/users). + + ``data`` holds ``base.User`` items. ``offset`` is the ``next`` value of the + previous page. Requires the manage_company permission. + """ + + return await self._list_page(base.User, "/v0/users", limit, offset) + + async def iter_users(self, limit: int = 0) -> AsyncIterator[base.User]: + """Yield up to ``limit`` of your company's users (0 for all of them).""" + + async for user in self._iter(self.users_page, limit): + yield user + async def get_users( self, limit: int | None = None, offset: int | str | None = None ) -> list[base.User] | base.ErrorResponse: - """List the users of your company (requires admin permissions).""" + """List your company's users, following the API's pagination. - return await self._model( - base.User, "get", "/v0/users", params={"limit": limit, "offset": offset} - ) + ``limit`` is the most users to return (None or 0 for all of them), and + ``offset`` starts after a previous page (its ``next`` value), see ``users_page``. + Requires the manage_company permission. + """ + + return await self._list(self.users_page, limit or 0, offset) async def get_user(self, user_id: str) -> base.User | base.ErrorResponse: """Get a user by ID (requires admin permissions).""" @@ -1505,9 +1566,14 @@ async def delete_user(self, user_id: str) -> None | base.ErrorResponse: return await self._empty("delete", f"/v0/users/{_seg(user_id)}") async def get_user_apikeys(self, user_id: str) -> list[base.ApiKey] | base.ErrorResponse: - """List the API keys of a user.""" + """List every API key of a user, following the API's pagination.""" + + uri = f"/v0/users/{_seg(user_id)}/apikeys" + + async def getpage(size, offset): + return await self._list_page(base.ApiKey, uri, size, offset) - return await self._model(base.ApiKey, "get", f"/v0/users/{_seg(user_id)}/apikeys") + return await self._list(getpage, 0) async def create_user_apikey(self, user_id: str, name: str) -> base.ApiKey | base.ErrorResponse: """Create a new API key for a user.""" diff --git a/src/pyhatching/base.py b/src/pyhatching/base.py index 0e3a7af..06b1260 100644 --- a/src/pyhatching/base.py +++ b/src/pyhatching/base.py @@ -84,6 +84,17 @@ class SamplesResponse(SampleInfo): """Response object for POST /samples (and the items of search results).""" +class ListPage(HatchingResponse): + """A single page of a paginated collection (ex. GET /profiles, GET /users). + + ``data`` holds the page's items, as the model the endpoint's method documents. + ``next`` is the offset to request the next page with, ``None`` on the last page. + """ + + data: list[Any] = Field(default_factory=list) + next: Optional[str] = None + + class SamplesPage(HatchingResponse): """A single page of samples, from GET /samples or GET /search. diff --git a/src/pyhatching/testing.py b/src/pyhatching/testing.py index 6bdc196..c1f6740 100644 --- a/src/pyhatching/testing.py +++ b/src/pyhatching/testing.py @@ -682,7 +682,8 @@ def _notfound(what: str = "sample") -> web.Response: @staticmethod def _page(request: web.Request, items: list) -> web.Response: - limit = int(request.query.get("limit", 20)) + # The API's default page size is 50, and its max is 200. + limit = min(int(request.query.get("limit", 50)), 200) offset = int(request.query.get("offset", 0) or 0) data = items[offset:offset + limit] nxt = str(offset + limit) if offset + limit < len(items) else None @@ -1008,7 +1009,7 @@ async def listKeys(self, request: web.Request) -> web.Response: uid = request.match_info["uid"] if uid not in self.users: return self._notfound("user") - return web.json_response({"data": self.apikeys.get(uid, [])}) + return self._page(request, self.apikeys.get(uid, [])) async def createKey(self, request: web.Request) -> web.Response: uid = request.match_info["uid"] diff --git a/tests/test_client.py b/tests/test_client.py index 4554a25..8538f44 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -400,6 +400,46 @@ async def test_profiles(client, fake): assert err.error == "NOT_FOUND" +async def test_profiles_pagination(client, fake): + for i in range(250): + fake.profiles[f"p{i}"] = {"id": f"p{i}", "name": f"p{i}", "tags": [], "network": ""} + + # Every page is followed, in pages of the max size. + profiles = await client.get_profiles() + assert len(profiles) == 251 and all(isinstance(p, base.HatchingProfileResponse) for p in profiles) + assert [r["query"] for r in fake.requests if r["path"].endswith("/profiles")] == [ + {"limit": "200"}, {"limit": "200", "offset": "200"}] + + assert [p.name for p in await client.get_profiles(limit=3)] == ["win10-internet", "p0", "p1"] + assert fake.last()["query"] == {"limit": "3"} + + page = await client.profiles_page(limit=100) + assert isinstance(page, base.ListPage) and len(page.data) == 100 and page.next == "100" + assert isinstance(page.data[0], base.HatchingProfileResponse) + page = await client.profiles_page(limit=500, offset=page.next) + assert len(page.data) == 151 and page.next is None + assert fake.last()["query"] == {"limit": "200", "offset": "100"} + + # Starting from a previous page's offset. + rest = await client.get_profiles(offset="250") + assert [p.name for p in rest] == ["p249"] + + names = [p.name async for p in client.iter_profiles(limit=201)] + assert len(names) == 201 and names[-1] == "p199" + assert len([p async for p in client.iter_profiles()]) == 251 + + +async def test_pagination_errors(fake): + async with pyhatching.PyHatchingClient("badkey", url=fake.url) as cli: + assert (await cli.get_profiles()).error == "UNAUTHORIZED" + assert (await cli.users_page()).error == "UNAUTHORIZED" + with pytest.raises(errors.PyHatchingApiError): + [u async for u in cli.iter_users()] + async with pyhatching.PyHatchingClient("badkey", url=fake.url, raise_on_api_err=True) as cli: + with pytest.raises(errors.PyHatchingApiError): + await cli.get_users() + + async def test_yara(client, fake): rules = await client.get_rules() assert rules.rules[0].name == "fake_rule.yara" @@ -442,10 +482,28 @@ async def test_users(client, fake): assert await client.delete_user_apikey(new.id, "ci") is None assert await client.get_user_apikeys(new.id) == [] + assert (await client.get_user_apikeys("nope")).error == "NOT_FOUND" + assert await client.delete_user(new.id) is None assert (await client.get_user(new.id)).error == "NOT_FOUND" +async def test_users_pagination(client, fake): + for i in range(60): + fake.users[f"user-x{i}"] = dict(ft.USER, id=f"user-x{i}", username=f"x{i}") + users = await client.get_users() + assert len(users) == 61 and isinstance(users[-1], base.User) + assert len(await client.get_users(limit=10)) == 10 + assert len([u async for u in client.iter_users(limit=0)]) == 61 + page = await client.users_page(limit=50) + assert len(page.data) == 50 and page.next == "50" + + fake.apikeys[ft.USER["id"]] = [{"name": f"k{i}", "key": "k" * 40} for i in range(210)] + keys = await client.get_user_apikeys(ft.USER["id"]) + assert len(keys) == 210 and isinstance(keys[0], base.ApiKey) + assert fake.last()["query"] == {"limit": "200", "offset": "200"} + + async def test_urlscan(client, fake): report = await client.urlscan_report(ft.SAMPLE_ID) assert report["url"] == "http://evil.example.com/login"