pangea-sdk 3.7.0__py3-none-any.whl → 3.8.0__py3-none-any.whl

pangea/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "3.7.0"
+__version__ = "3.8.0"
 
 from pangea.asyncio.request import PangeaRequestAsync
 from pangea.config import PangeaConfig

pangea/asyncio/request.py

@@ -7,14 +7,16 @@ import time
 from typing import Dict, List, Optional, Tuple, Type, Union
 
 import aiohttp
-import pangea.exceptions as pe
 from aiohttp import FormData
+from typing_extensions import TypeVar
 
-# from requests.adapters import HTTPAdapter, Retry
-from pangea.request import PangeaRequestBase
-from pangea.response import AcceptedResult, PangeaResponse, PangeaResponseResult, ResponseStatus, TransferMethod
+import pangea.exceptions as pe
+from pangea.request import MultipartResponse, PangeaRequestBase
+from pangea.response import AttachedFile, PangeaResponse, PangeaResponseResult, ResponseStatus, TransferMethod
 from pangea.utils import default_encoder
 
+TResult = TypeVar("TResult", bound=PangeaResponseResult, default=PangeaResponseResult)
+
 
 class PangeaRequestAsync(PangeaRequestBase):
     """An object that makes direct calls to Pangea Service APIs.
@@ -28,12 +30,12 @@ class PangeaRequestAsync(PangeaRequestBase):
     async def post(
         self,
         endpoint: str,
-        result_class: Type[PangeaResponseResult],
+        result_class: Type[TResult],
         data: Union[str, Dict] = {},
-        files: Optional[List[Tuple]] = None,
+        files: List[Tuple] = [],
         poll_result: bool = True,
         url: Optional[str] = None,
-    ) -> PangeaResponse:
+    ) -> PangeaResponse[TResult]:
         """Makes the POST call to a Pangea Service endpoint.
 
         Args:
@@ -56,7 +58,7 @@ class PangeaRequestAsync(PangeaRequestBase):
         )
         transfer_method = data.get("transfer_method", None)  # type: ignore[union-attr]
 
-        if files is not None and type(data) is dict and (transfer_method == TransferMethod.POST_URL.value):
+        if files and type(data) is dict and (transfer_method == TransferMethod.POST_URL.value):
             requests_response = await self._full_post_presigned_url(
                 endpoint, result_class=result_class, data=data, files=files
             )
@@ -66,18 +68,32 @@ class PangeaRequestAsync(PangeaRequestBase):
             )
 
         await self._check_http_errors(requests_response)
-        json_resp = await requests_response.json()
-        self.logger.debug(json.dumps({"service": self.service, "action": "post", "url": url, "response": json_resp}))
 
-        pangea_response = PangeaResponse(requests_response, result_class=result_class, json=json_resp)  # type: ignore[var-annotated]
+        if "multipart/form-data" in requests_response.headers.get("content-type", ""):
+            multipart_response = await self._process_multipart_response(requests_response)
+            pangea_response: PangeaResponse = PangeaResponse(
+                requests_response,
+                result_class=result_class,
+                json=multipart_response.pangea_json,
+                attached_files=multipart_response.attached_files,
+            )
+        else:
+            try:
+                json_resp = await requests_response.json()
+                self.logger.debug(
+                    json.dumps({"service": self.service, "action": "post", "url": url, "response": json_resp})
+                )
+
+                pangea_response = PangeaResponse(requests_response, result_class=result_class, json=json_resp)
+            except aiohttp.ContentTypeError as e:
+                raise pe.PangeaException(f"Failed to decode json response. {e}. Body: {await requests_response.text()}")
+
         if poll_result:
             pangea_response = await self._handle_queued_result(pangea_response)
 
         return self._check_response(pangea_response)
 
-    async def get(
-        self, path: str, result_class: Type[PangeaResponseResult], check_response: bool = True
-    ) -> PangeaResponse:
+    async def get(self, path: str, result_class: Type[TResult], check_response: bool = True) -> PangeaResponse[TResult]:
         """Makes the GET call to a Pangea Service endpoint.
 
         Args:
@@ -94,7 +110,7 @@ class PangeaRequestAsync(PangeaRequestBase):
 
         async with self.session.get(url, headers=self._headers()) as requests_response:
             await self._check_http_errors(requests_response)
-            pangea_response = PangeaResponse(  # type: ignore[var-annotated]
+            pangea_response = PangeaResponse(
                 requests_response, result_class=result_class, json=await requests_response.json()
             )
 
@@ -115,11 +131,11 @@ class PangeaRequestAsync(PangeaRequestBase):
             raise pe.ServiceTemporarilyUnavailable(await resp.json())
 
     async def poll_result_by_id(
-        self, request_id: str, result_class: Union[Type[PangeaResponseResult], dict], check_response: bool = True
-    ):
+        self, request_id: str, result_class: Type[TResult], check_response: bool = True
+    ) -> PangeaResponse[TResult]:
         path = self._get_poll_path(request_id)
         self.logger.debug(json.dumps({"service": self.service, "action": "poll_result_once", "url": path}))
-        return await self.get(path, result_class, check_response=check_response)  # type: ignore[arg-type]
+        return await self.get(path, result_class, check_response=check_response)
 
     async def poll_result_once(self, response: PangeaResponse, check_response: bool = True):
         request_id = response.request_id
@@ -157,12 +173,84 @@ class PangeaRequestAsync(PangeaRequestBase):
         if resp.status < 200 or resp.status >= 300:
             raise pe.PresignedUploadError(f"presigned PUT failure: {resp.status}", await resp.text())
 
+    async def download_file(self, url: str, filename: Optional[str] = None) -> AttachedFile:
+        self.logger.debug(
+            json.dumps(
+                {
+                    "service": self.service,
+                    "action": "download_file",
+                    "url": url,
+                    "status": "start",
+                }
+            )
+        )
+        async with self.session.get(url, headers={}) as response:
+            if response.status == 200:
+                if filename is None:
+                    content_disposition = response.headers.get("Content-Disposition", "")
+                    filename = self._get_filename_from_content_disposition(content_disposition)
+                    if filename is None:
+                        filename = self._get_filename_from_url(url)
+                        if filename is None:
+                            filename = "default_filename"
+
+                content_type = response.headers.get("Content-Type", "")
+                self.logger.debug(
+                    json.dumps(
+                        {
+                            "service": self.service,
+                            "action": "download_file",
+                            "url": url,
+                            "filename": filename,
+                            "status": "success",
+                        }
+                    )
+                )
+
+                return AttachedFile(filename=filename, file=await response.read(), content_type=content_type)
+            else:
+                raise pe.DownloadFileError(f"Failed to download file. Status: {response.status}", await response.text())
+
+    async def _get_pangea_json(self, reader: aiohttp.MultipartReader) -> Optional[Dict]:
+        # Iterate through parts
+        async for part in reader:
+            return await part.json()
+        return None
+
+    async def _get_attached_files(self, reader: aiohttp.MultipartReader) -> List[AttachedFile]:
+        files = []
+        i = 0
+
+        async for part in reader:
+            content_type = part.headers.get("Content-Type", "")
+            content_disposition = part.headers.get("Content-Disposition", "")
+            name = self._get_filename_from_content_disposition(content_disposition)
+            if name is None:
+                name = f"default_file_name_{i}"
+                i += 1
+            files.append(AttachedFile(name, await part.read(), content_type))
+
+        return files
+
+    async def _process_multipart_response(self, resp: aiohttp.ClientResponse) -> MultipartResponse:
+        # Parse the multipart response
+        multipart_reader = aiohttp.MultipartReader.from_response(resp)
+
+        pangea_json = await self._get_pangea_json(multipart_reader)  # type: ignore[arg-type]
+        self.logger.debug(
+            json.dumps({"service": self.service, "action": "multipart response", "response": pangea_json})
+        )
+
+        multipart_reader = multipart_reader.__aiter__()
+        attached_files = await self._get_attached_files(multipart_reader)  # type: ignore[arg-type]
+        return MultipartResponse(pangea_json, attached_files)  # type: ignore[arg-type]
+
     async def _http_post(
         self,
         url: str,
         headers: Dict = {},
         data: Union[str, Dict] = {},
-        files: Optional[List[Tuple]] = None,
+        files: List[Tuple] = [],
         presigned_url_post: bool = False,
     ) -> aiohttp.ClientResponse:
         if files:
@@ -193,26 +281,29 @@ class PangeaRequestAsync(PangeaRequestBase):
         self.logger.debug(
             json.dumps({"service": self.service, "action": "http_put", "url": url}, default=default_encoder)
         )
-        form = FormData()
-        name, value = files[0]
-        form.add_field(name, value[1], filename=value[0], content_type=value[2])
-        return await self.session.put(url, headers=headers, data=form)
+        _, value = files[0]
+        return await self.session.put(url, headers=headers, data=value[1])
 
     async def _full_post_presigned_url(
         self,
         endpoint: str,
         result_class: Type[PangeaResponseResult],
         data: Union[str, Dict] = {},
-        files: Optional[List[Tuple]] = None,
+        files: List[Tuple] = [],
     ):
-        if len(files) == 0:  # type: ignore[arg-type]
+        if len(files) == 0:
             raise AttributeError("files attribute should have at least 1 file")
 
         response = await self.request_presigned_url(endpoint=endpoint, result_class=result_class, data=data)
-        data_to_presigned = response.accepted_result.post_form_data  # type: ignore[union-attr]
-        presigned_url = response.accepted_result.post_url  # type: ignore[union-attr]
+        if response.accepted_result is None:
+            raise pe.PangeaException("No accepted_result field when requesting presigned url")
+        if response.accepted_result.post_url is None:
+            raise pe.PresignedURLException("No presigned url", response)
+
+        data_to_presigned = response.accepted_result.post_form_data
+        presigned_url = response.accepted_result.post_url
 
-        await self.post_presigned_url(url=presigned_url, data=data_to_presigned, files=files)  # type: ignore[arg-type]
+        await self.post_presigned_url(url=presigned_url, data=data_to_presigned, files=files)
         return response.raw_response
 
     async def request_presigned_url(
@@ -232,14 +323,14 @@ class PangeaRequestAsync(PangeaRequestBase):
             raise e
 
         # Receive 202
-        return await self._poll_presigned_url(accepted_exception.response)  # type: ignore[return-value]
+        return await self._poll_presigned_url(accepted_exception.response)
 
-    async def _poll_presigned_url(self, response: PangeaResponse) -> AcceptedResult:
+    async def _poll_presigned_url(self, response: PangeaResponse[TResult]) -> PangeaResponse[TResult]:
         if response.http_status != 202:
             raise AttributeError("Response should be 202")
 
         if response.accepted_result is not None and response.accepted_result.has_upload_url:
-            return response  # type: ignore[return-value]
+            return response
 
         self.logger.debug(json.dumps({"service": self.service, "action": "poll_presigned_url", "step": "start"}))
         retry_count = 1
@@ -276,7 +367,7 @@ class PangeaRequestAsync(PangeaRequestBase):
         self.logger.debug(json.dumps({"service": self.service, "action": "poll_presigned_url", "step": "exit"}))
 
         if loop_resp.accepted_result is not None and not loop_resp.accepted_result.has_upload_url:
-            return loop_resp  # type: ignore[return-value]
+            return loop_resp
         else:
             raise loop_exc
 

pangea/asyncio/services/__init__.py

@@ -1,5 +1,6 @@
 from .audit import AuditAsync
 from .authn import AuthNAsync
+from .authz import AuthZAsync
 from .embargo import EmbargoAsync
 from .file_scan import FileScanAsync
 from .intel import DomainIntelAsync, FileIntelAsync, IpIntelAsync, UrlIntelAsync, UserIntelAsync

pangea/asyncio/services/audit.py

@@ -1,18 +1,28 @@
 # Copyright 2022 Pangea Cyber Corporation
 # Author: Pangea Cyber Corporation
+from __future__ import annotations
+
 import datetime
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 import pangea.exceptions as pexc
-from pangea.response import PangeaResponse
+from pangea.asyncio.services.base import ServiceBaseAsync
+from pangea.config import PangeaConfig
+from pangea.response import PangeaResponse, PangeaResponseResult
 from pangea.services.audit.audit import AuditBase
 from pangea.services.audit.exceptions import AuditException
 from pangea.services.audit.models import (
+    DownloadFormat,
+    DownloadRequest,
+    DownloadResult,
     Event,
+    ExportRequest,
     LogBulkResult,
     LogResult,
+    PublishedRoot,
     RootRequest,
     RootResult,
+    RootSource,
     SearchOrder,
     SearchOrderBy,
     SearchOutput,
@@ -22,8 +32,6 @@ from pangea.services.audit.models import (
 )
 from pangea.services.audit.util import format_datetime
 
-from .base import ServiceBaseAsync
-
 
 class AuditAsync(ServiceBaseAsync, AuditBase):
     """Audit service client.
@@ -52,14 +60,33 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
 
     def __init__(
         self,
-        token,
-        config=None,
+        token: str,
+        config: PangeaConfig | None = None,
         private_key_file: str = "",
-        public_key_info: Dict[str, str] = {},
-        tenant_id: Optional[str] = None,
-        logger_name="pangea",
-        config_id: Optional[str] = None,
-    ):
+        public_key_info: dict[str, str] = {},
+        tenant_id: str | None = None,
+        logger_name: str = "pangea",
+        config_id: str | None = None,
+    ) -> None:
+        """
+        Audit client
+
+        Initializes a new Audit client.
+
+        Args:
+            token: Pangea API token.
+            config: Configuration.
+            private_key_file: Private key filepath.
+            public_key_info: Public key information.
+            tenant_id: Tenant ID.
+            logger_name: Logger name.
+            config_id: Configuration ID.
+
+        Examples:
+             config = PangeaConfig(domain="pangea_domain")
+             audit = AuditAsync(token="pangea_token", config=config)
+        """
+
         # FIXME: Temporary check to deprecate config_id from PangeaConfig.
         # Delete it when deprecate PangeaConfig.config_id
         if config_id and config is not None and config.config_id is not None:
@@ -175,8 +202,10 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
         """
 
         input = self._get_log_request(event, sign_local=sign_local, verify=verify, verbose=verbose)
-        response = await self.request.post("v1/log", LogResult, data=input.dict(exclude_none=True))
-        if response.success:
+        response: PangeaResponse[LogResult] = await self.request.post(
+            "v1/log", LogResult, data=input.dict(exclude_none=True)
+        )
+        if response.success and response.result is not None:
             self._process_log_result(response.result, verify=verify)
         return response
 
@@ -209,8 +238,10 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
         """
 
         input = self._get_log_request(events, sign_local=sign_local, verify=False, verbose=verbose)
-        response = await self.request.post("v2/log", LogBulkResult, data=input.dict(exclude_none=True))
-        if response.success:
+        response: PangeaResponse[LogBulkResult] = await self.request.post(
+            "v2/log", LogBulkResult, data=input.dict(exclude_none=True)
+        )
+        if response.success and response.result is not None:
             for result in response.result.results:
                 self._process_log_result(result, verify=True)
         return response
@@ -245,12 +276,12 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
 
         input = self._get_log_request(events, sign_local=sign_local, verify=False, verbose=verbose)
         try:
-            response = await self.request.post(
+            response: PangeaResponse[LogBulkResult] = await self.request.post(
                 "v2/log_async", LogBulkResult, data=input.dict(exclude_none=True), poll_result=False
             )
         except pexc.AcceptedRequestException as e:
             return e.response
-        if response.success:
+        if response.success and response.result:
             for result in response.result.results:
                 self._process_log_result(result, verify=True)
         return response
@@ -264,7 +295,7 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
         end: Optional[Union[datetime.datetime, str]] = None,
         limit: Optional[int] = None,
         max_results: Optional[int] = None,
-        search_restriction: Optional[dict] = None,
+        search_restriction: Optional[Dict[str, Sequence[str]]] = None,
         verbose: Optional[bool] = None,
         verify_consistency: bool = False,
         verify_events: bool = True,
@@ -293,7 +324,7 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
             end (datetime, optional): An RFC-3339 formatted timestamp, or relative time adjustment from the current time.
             limit (int, optional): Optional[int] = None,
             max_results (int, optional): Maximum number of results to return.
-            search_restriction (dict, optional): A list of keys to restrict the search results to. Useful for partitioning data available to the query string.
+            search_restriction (Dict[str, Sequence[str]], optional): A list of keys to restrict the search results to. Useful for partitioning data available to the query string.
             verbose (bool, optional): If true, response include root and membership and consistency proofs.
             verify_consistency (bool): True to verify logs consistency
             verify_events (bool): True to verify hash events and signatures
@@ -326,7 +357,12 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
             verbose=verbose,
         )
 
-        response = await self.request.post("v1/search", SearchOutput, data=input.dict(exclude_none=True))
+        response: PangeaResponse[SearchOutput] = await self.request.post(
+            "v1/search", SearchOutput, data=input.dict(exclude_none=True)
+        )
+        if verify_consistency:
+            await self.update_published_roots(response.result)  # type: ignore[arg-type]
+
         return self.handle_search_response(response, verify_consistency, verify_events)
 
     async def results(
@@ -334,6 +370,7 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
         id: str,
         limit: Optional[int] = 20,
         offset: Optional[int] = 0,
+        assert_search_restriction: Optional[Dict[str, Sequence[str]]] = None,
         verify_consistency: bool = False,
         verify_events: bool = True,
     ) -> PangeaResponse[SearchResultOutput]:
@@ -348,6 +385,7 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
             id (string): the id of a search action, found in `response.result.id`
             limit (integer, optional): the maximum number of results to return, default is 20
             offset (integer, optional): the position of the first result to return, default is 0
+            assert_search_restriction (Dict[str, Sequence[str]], optional): Assert the requested search results were queried with the exact same search restrictions, to ensure the results comply to the expected restrictions.
             verify_consistency (bool): True to verify logs consistency
             verify_events (bool): True to verify hash events and signatures
         Raises:
@@ -365,7 +403,8 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
             result_res: PangeaResponse[SearchResultsOutput] = audit.results(
                 id=search_res.result.id,
                 limit=10,
-                offset=0)
+                offset=0,
+                assert_search_restriction={'source': ["monitor"]})
         """
 
         if limit <= 0:  # type: ignore[operator]
@@ -378,10 +417,115 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
             id=id,
             limit=limit,
             offset=offset,
+            assert_search_restriction=assert_search_restriction,
         )
         response = await self.request.post("v1/results", SearchResultOutput, data=input.dict(exclude_none=True))
+        if verify_consistency and response.result is not None:
+            await self.update_published_roots(response.result)
+
         return self.handle_results_response(response, verify_consistency, verify_events)
 
+    async def export(
+        self,
+        *,
+        format: DownloadFormat = DownloadFormat.CSV,
+        start: Optional[datetime.datetime] = None,
+        end: Optional[datetime.datetime] = None,
+        order: Optional[SearchOrder] = None,
+        order_by: Optional[str] = None,
+        verbose: bool = True,
+    ) -> PangeaResponse[PangeaResponseResult]:
+        """
+        Export from the audit log
+
+        Bulk export of data from the Secure Audit Log, with optional filtering.
+
+        OperationId: audit_post_v1_export
+
+        Args:
+            format: Format for the records.
+            start: The start of the time range to perform the search on.
+            end: The end of the time range to perform the search on. If omitted,
+                then all records up to the latest will be searched.
+            order: Specify the sort order of the response.
+            order_by: Name of column to sort the results by.
+            verbose: Whether or not to include the root hash of the tree and the
+                membership proof for each record.
+
+        Raises:
+            AuditException: If an audit based api exception happens
+            PangeaAPIException: If an API Error happens
+
+        Examples:
+            export_res = await audit.export(verbose=False)
+
+            # Export may take several dozens of minutes, so polling for the result
+            # should be done in a loop. That is omitted here for brevity's sake.
+            try:
+                await audit.poll_result(request_id=export_res.request_id)
+            except AcceptedRequestException:
+                # Retry later.
+
+            # Download the result when it's ready.
+            download_res = await audit.download_results(request_id=export_res.request_id)
+            download_res.result.dest_url
+            # => https://pangea-runtime.s3.amazonaws.com/audit/xxxxx/search_results_[...]
+        """
+        input = ExportRequest(
+            format=format,
+            start=start,
+            end=end,
+            order=order,
+            order_by=order_by,
+            verbose=verbose,
+        )
+        try:
+            return await self.request.post(
+                "v1/export", PangeaResponseResult, data=input.dict(exclude_none=True), poll_result=False
+            )
+        except pexc.AcceptedRequestException as e:
+            return e.response
+
+    async def log_stream(self, data: dict) -> PangeaResponse[PangeaResponseResult]:
+        """
+        Log streaming endpoint
+
+        This API allows 3rd party vendors (like Auth0) to stream events to this
+        endpoint where the structure of the payload varies across different
+        vendors.
+
+        OperationId: audit_post_v1_log_stream
+
+        Args:
+            data: Event data. The exact schema of this will vary by vendor.
+
+        Raises:
+            AuditException: If an audit based api exception happens
+            PangeaAPIException: If an API Error happens
+
+        Examples:
+            data = {
+                "logs": [
+                    {
+                        "log_id": "some log ID",
+                        "data": {
+                            "date": "2024-03-29T17:26:50.193Z",
+                            "type": "sapi",
+                            "description": "Create a log stream",
+                            "client_id": "some client ID",
+                            "ip": "127.0.0.1",
+                            "user_agent": "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0",
+                            "user_id": "some user ID",
+                        },
+                    }
+                    # ...
+                ]
+            }
+
+            response = await audit.log_stream(data)
+        """
+        return await self.request.post("v1/log_stream", PangeaResponseResult, data=data)
+
     async def root(self, tree_size: Optional[int] = None) -> PangeaResponse[RootResult]:
         """
         Tamperproof verification
@@ -405,3 +549,74 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
         """
         input = RootRequest(tree_size=tree_size)
         return await self.request.post("v1/root", RootResult, data=input.dict(exclude_none=True))
+
+    async def download_results(
+        self,
+        result_id: Optional[str] = None,
+        format: DownloadFormat = DownloadFormat.CSV,
+        request_id: Optional[str] = None,
+    ) -> PangeaResponse[DownloadResult]:
+        """
+        Download search results
+
+        Get all search results as a compressed (gzip) CSV file.
+
+        OperationId: audit_post_v1_download_results
+
+        Args:
+            result_id: ID returned by the search API.
+            format: Format for the records.
+            request_id: ID returned by the export API.
+
+        Returns:
+            URL where search results can be downloaded.
+
+        Raises:
+            AuditException: If an Audit-based API exception occurs.
+            PangeaAPIException: If an API exception occurs.
+
+        Examples:
+            response = await audit.download_results(
+                result_id="pas_[...]",
+                format=DownloadFormat.JSON,
+            )
+        """
+
+        if request_id is None and result_id is None:
+            raise ValueError("must pass one of `request_id` or `result_id`")
+
+        input = DownloadRequest(request_id=request_id, result_id=result_id, format=format)
+        return await self.request.post("v1/download_results", DownloadResult, data=input.dict(exclude_none=True))
+
+    async def update_published_roots(self, result: SearchResultOutput):
+        """Fetches series of published root hashes from Arweave
+
+        This is used for subsequent calls to verify_consistency_proof(). Root hashes
+        are published on [Arweave](https://arweave.net).
+
+        Args:
+            result (SearchResultOutput): Result object from previous call to AuditAsync.search() or AuditAsync.results()
+
+        Raises:
+            AuditException: If an audit based api exception happens
+            PangeaAPIException: If an API Error happens
+        """
+
+        if not result.root:
+            return
+
+        tree_sizes, arweave_roots = self._get_tree_sizes_and_roots(result)
+
+        # fill the missing roots from the server (if allowed)
+        for tree_size in tree_sizes:
+            pub_root = None
+            if tree_size in arweave_roots:
+                pub_root = PublishedRoot(**arweave_roots[tree_size].dict(exclude_none=True))
+                pub_root.source = RootSource.ARWEAVE
+            elif self.allow_server_roots:
+                resp = await self.root(tree_size=tree_size)
+                if resp.success and resp.result is not None:
+                    pub_root = PublishedRoot(**resp.result.data.dict(exclude_none=True))
+                    pub_root.source = RootSource.PANGEA
+            if pub_root is not None:
+                self.pub_roots[tree_size] = pub_root