Merge pull request #159 from censys/adh/comments-and-tags

feat(api): Add Tagging and Comments

Author: thehappydinoa

.github/workflows/python-ci.yml

@@ -35,7 +35,6 @@ jobs:
           key: venv-${{ runner.os }}-${{ hashFiles('**/poetry.lock') }}
 
       - name: Install dependencies
-        if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
         run: |
           poetry install
 

censys/common/base.py

@@ -3,7 +3,7 @@
 import os
 import warnings
 from functools import wraps
-from typing import Any, Callable, List, Optional, Type
+from typing import Any, Callable, List, Optional, Protocol, Type
 
 import backoff
 import requests
@@ -43,6 +43,11 @@ def _impl():
     return _wrapper
 
 
+class RequestsMethod(Protocol):  # noqa: D101
+    def __call__(self, path: str, **kwargs) -> requests.Response:  # noqa: D102
+        ...  # pragma: no cover
+
+
 class CensysAPIBase:
     """This is the base class for API queries."""
 
@@ -126,7 +131,7 @@ def _get_exception_class(_: Response) -> Type[CensysAPIException]:
     @_backoff_wrapper
     def _make_call(
         self,
-        method: Callable,
+        method: RequestsMethod,
         endpoint: str,
         args: Optional[dict] = None,
         data: Optional[Any] = None,
@@ -137,7 +142,7 @@ def _make_call(
         and decoding the responses.
 
         Args:
-            method (Callable): Method to send HTTP request.
+            method (RequestsMethod): Method to send HTTP request.
             endpoint (str): The path of API endpoint.
             args (dict): Optional; URL args that are mapped to params.
             data (Any): Optional; JSON data to serialize with request.
@@ -164,15 +169,18 @@ def _make_call(
 
         res = method(url, **request_kwargs)
 
-        if res.status_code == 200:
+        if res.status_code in range(200, 300):
             # Check for a returned json body
             try:
                 json_data = res.json()
                 if "error" not in json_data:
                     return json_data
             # Successful request returned no json body in response
             except ValueError:
-                return {}
+                return {
+                    "code": res.status_code,
+                    "status": res.reason,
+                }
 
         try:
             json_data = res.json()

censys/search/v2/api.py

@@ -1,7 +1,7 @@
 """Base for interacting with the Censys Search API."""
 import os
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import Dict, Iterable, Iterator, List, Optional, Type
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Type
 
 from requests.models import Response
 
@@ -74,6 +74,7 @@ def __init__(
         self.search_path = f"/{self.INDEX_NAME}/search"
         self.aggregate_path = f"/{self.INDEX_NAME}/aggregate"
         self.metadata_path = f"/metadata/{self.INDEX_NAME}"
+        self.tags_path = "/tags"
 
         # Set up the v1 API
         v1_kwargs = kwargs.copy()
@@ -286,3 +287,141 @@ def metadata(self) -> dict:
             dict: The result set returned.
         """
         return self._get(self.metadata_path)["result"]
+
+    # Comments
+
+    def get_comments(self, document_id: str) -> List[dict]:
+        """Get comments for a document.
+
+        Args:
+            document_id (str): The ID of the document you are requesting.
+
+        Returns:
+            List[dict]: The list of comments.
+        """
+        return self._get(self.view_path + document_id + "/comments")["result"][
+            "comments"
+        ]
+
+    def add_comment(self, document_id: str, contents: str) -> dict:
+        """Add comment to a document.
+
+        Args:
+            document_id (str): The ID of the document you are requesting.
+            contents (str): The contents of the comment.
+
+        Returns:
+            dict: The result set returned.
+        """
+        return self._post(
+            self.view_path + document_id + "/comments", data={"contents": contents}
+        )["result"]
+
+    # Tags
+
+    def list_all_tags(self) -> List[dict]:
+        """List all tags.
+
+        Returns:
+            List[dict]: The list of tags.
+        """
+        return self._get(self.tags_path)["result"]["tags"]
+
+    def create_tag(self, name: str, color: Optional[str] = None) -> dict:
+        """Create a tag.
+
+        Args:
+            name (str): The name of the tag.
+            color (str): Optional; The color of the tag.
+
+        Returns:
+            dict: The result set returned.
+        """
+        tag_def: Dict[str, Any] = {"name": name}
+        if color:
+            tag_def["metadata"] = {"color": color}
+        return self._post(self.tags_path, data=tag_def)["result"]
+
+    def get_tag(self, tag_id: str) -> dict:
+        """Get a tag.
+
+        Args:
+            tag_id (str): The ID of the tag.
+
+        Returns:
+            dict: The result set returned.
+        """
+        return self._get(self.tags_path + "/" + tag_id)["result"]
+
+    def update_tag(self, tag_id: str, name: str, color: Optional[str] = None) -> dict:
+        """Update a tag.
+
+        Args:
+            tag_id (str): The ID of the tag.
+            name (str): The name of the tag.
+            color (str): The color of the tag.
+
+        Returns:
+            dict: The result set returned.
+        """
+        tag_def: Dict[str, Any] = {"name": name}
+        if color:
+            tag_def["metadata"] = {"color": color}
+        return self._put(
+            self.tags_path + "/" + tag_id,
+            data=tag_def,
+        )["result"]
+
+    def delete_tag(self, tag_id: str):
+        """Delete a tag.
+
+        Args:
+            tag_id (str): The ID of the tag.
+        """
+        self._delete(self.tags_path + "/" + tag_id)
+
+    def _list_documents_with_tag(
+        self, tag_id: str, endpoint: str, keyword: str
+    ) -> List[dict]:
+        """List documents by tag.
+
+        Args:
+            tag_id (str): The ID of the tag.
+            endpoint (str): The endpoint to be called.
+            keyword (str): The keyword to be used in the endpoint.
+
+        Returns:
+            List[dict]: The list of documents.
+        """
+        return self._get(self.tags_path + "/" + tag_id + "/" + endpoint)["result"][
+            keyword
+        ]
+
+    def list_tags_on_document(self, document_id: str) -> List[dict]:
+        """List tags on a document.
+
+        Args:
+            document_id (str): The ID of the document.
+
+        Returns:
+            List[dict]: The list of tags.
+        """
+        return self._get(self.view_path + document_id + "/tags")["result"]["tags"]
+
+    def add_tag_to_document(self, document_id: str, tag_id: str):
+        """Add a tag to a document.
+
+        Args:
+            document_id (str): The ID of the document.
+            tag_id (str): The ID of the tag.
+        """
+        self._put(self.view_path + document_id + "/tags/" + tag_id)
+
+    def remove_tag_from_document(self, document_id: str, tag_id: str):
+        """Remove a tag from a document.
+
+        Args:
+            document_id (str): The ID of the document.
+            tag_id (str): The ID of the tag.
+        """
+        self._delete(self.view_path + document_id + "/tags/" + tag_id)

censys/search/v2/certs.py

@@ -1,4 +1,4 @@
-"""Interact with the Censys Search Host API."""
+"""Interact with the Censys Search Cert API."""
 from typing import List, Optional, Tuple
 
 from .api import CensysSearchAPIv2
@@ -16,14 +16,19 @@ class CensysCerts(CensysSearchAPIv2):
         Search for hosts by sha256fp.
 
         >>> c.get_hosts_by_cert("fb444eb8e68437bae06232b9f5091bccff62a768ca09e92eb5c9c2cf9d17c426")
-        [
+        (
+            [
+                {
+                    "ip": "string",
+                    "name": "string",
+                    "observed_at": "2021-08-02T14:56:38.711Z",
+                    "first_observed_at": "2021-08-02T14:56:38.711Z",
+                }
+            ],
             {
-                "ip": "string",
-                "name": "string",
-                "observed_at": "2021-08-02T14:56:38.711Z",
-                "first_observed_at": "2021-08-02T14:56:38.711Z",
-            }
-        ]
+                "next": "nextCursorToken",
+            },
+        )
     """
 
     INDEX_NAME = "certificates"
@@ -58,3 +63,15 @@ def get_hosts_by_cert(
         args = {"cursor": cursor}
         result = self._get(self.view_path + sha256fp + "/hosts", args)["result"]
         return result["hosts"], result["links"]
+
+    def list_certs_with_tag(self, tag_id: str) -> List[str]:
+        """Returns a list of certs which are tagged with the specified tag.
+
+        Args:
+            tag_id (str): The ID of the tag.
+
+        Returns:
+            List[str]: A list of certificate SHA 256 fingerprints.
+        """
+        certs = self._list_documents_with_tag(tag_id, "certificates", "certs")
+        return [cert["fingerprint"] for cert in certs]

censys/search/v2/hosts.py

@@ -31,7 +31,7 @@ class CensysHosts(CensysSearchAPIv2):
             ...
         ]
 
-        View specific host.
+        Fetch a specific host and its services
 
         >>> h.view("1.0.0.0")
         {
@@ -57,6 +57,16 @@ class CensysHosts(CensysSearchAPIv2):
             'query': 'service.service_name: HTTP',
             'total': 172588754
         }
+
+        Fetch a list of host names for the specified IP address.
+
+        >>> h.view_host_names("1.1.1.1")
+        ['one.one.one.one']
+
+        Fetch a list of events for the specified IP address.
+
+        >>> h.view_host_events("1.1.1.1")
+        [{'timestamp': '2019-01-01T00:00:00.000Z'}]
     """
 
     INDEX_NAME = "hosts"
@@ -108,3 +118,15 @@ def view_host_events(
         return self._get(f"/experimental/{self.INDEX_NAME}/{ip_address}/events", args)[
             "result"
         ]["events"]
+
+    def list_hosts_with_tag(self, tag_id: str) -> List[str]:
+        """Returns a list of hosts which are tagged with the specified tag.
+
+        Args:
+            tag_id (str): The ID of the tag.
+
+        Returns:
+            List[str]: A list of host IP addresses.
+        """
+        hosts = self._list_documents_with_tag(tag_id, "hosts", "hosts")
+        return [host["ip"] for host in hosts]

docs/usage-v1.rst

@@ -132,7 +132,7 @@ Below we show an example using the :attr:`CensysIPv4 <censys.search.v1.CensysIPv
 ``bulk``
 --------
 
-**Please note this method is only available only for the certificate index**
+**Please note this method is only available only for the CensysCertificates index**
 
 Below we show an example using the :attr:`CensysCertificates <censys.search.v1.CensysCertificates>` index.