feat: Add CLI and API support for ASM saved queries (#598)

Author: rickydia

Diff for: README.md

@@ -20,7 +20,9 @@ An easy-to-use and lightweight API wrapper for Censys APIs ([censys.io](https://
 - [Command-line interface](https://censys-python.readthedocs.io/en/stable/cli.html)
 
 <!-- markdownlint-disable MD033 -->
+
 <a href="https://asciinema.org/a/500416" target="_blank"><img src="https://asciinema.org/a/500416.svg" width="600"/></a>
+
 <!-- markdownlint-enable MD033 -->
 
 ## Getting Started
@@ -121,4 +123,4 @@ poetry run pytest --cov-report html
 
 This software is licensed under [Apache License, Version 2.0](http://www.apache.org/licenses/LICENSE-2.0)
 
-- Copyright (C) 2023 Censys, Inc.
+- Copyright (C) 2024 Censys, Inc.

Diff for: censys/asm/__init__.py

@@ -14,6 +14,7 @@
 from .inventory import InventorySearch
 from .logbook import Events, Logbook
 from .risks import Risks
+from .saved_queries import SavedQueries
 from .seeds import Seeds
 
 __all__ = [
@@ -28,6 +29,7 @@
     "InventorySearch",
     "Logbook",
     "Risks",
+    "SavedQueries",
     "Seeds",
     "SubdomainsAssets",
     "WebEntitiesAssets",

Diff for: censys/asm/client.py

@@ -14,6 +14,7 @@
 from .inventory import InventorySearch
 from .logbook import Logbook
 from .risks import Risks
+from .saved_queries import SavedQueries
 from .seeds import Seeds
 
 
@@ -40,3 +41,4 @@ def __init__(self, api_key: Optional[str] = None, **kwargs):
         self.object_storages = ObjectStoragesAssets(api_key, **kwargs)
         self.web_entities = WebEntitiesAssets(api_key, **kwargs)
         self.beta = Beta(api_key, **kwargs)
+        self.saved_queries = SavedQueries(api_key, **kwargs)

Diff for: censys/asm/saved_queries.py

@@ -0,0 +1,116 @@
+"""Interact with the Censys Saved Queries API."""
+from typing import Optional
+
+from .api import CensysAsmAPI
+
+
+class SavedQueries(CensysAsmAPI):
+    """Saved Queries API class."""
+
+    base_path = "/inventory/v1/saved-query"
+
+    def get_saved_queries(
+        self,
+        query_name_prefix: Optional[str] = None,
+        page_size: Optional[int] = None,
+        page: Optional[int] = None,
+        filter_term: Optional[str] = None,
+    ) -> dict:
+        """Get saved queries.
+
+        Args:
+            query_name_prefix (str, optional): Prefix for the saved query name.
+            page_size (int, optional): Number of results to return. Defaults to 50.
+            page (int, optional): Page number to begin at when searching. Defaults to 1.
+            filter_term (str, optional): Term used to filter the list of saved query names and the saved queries.
+
+        Returns:
+            dict: Saved queries results.
+        """
+        if page_size is None:
+            page_size = 50
+        if page is None:
+            page = 1
+        args: dict = {
+            "pageSize": page_size,
+            "page": page,
+        }
+
+        if query_name_prefix:
+            args["queryNamePrefix"] = query_name_prefix
+        if filter_term:
+            args["filterTerm"] = filter_term
+
+        return self._get(self.base_path, args=args)
+
+    def add_saved_query(
+        self,
+        query: str,
+        query_name: str,
+    ) -> dict:
+        """Add a new saved query to the ASM platform.
+
+        Args:
+            query (str): Query string.
+            query_name (str): Saved query name.
+
+        Returns:
+            dict: Added saved query results.
+        """
+        body = {
+            "query": query,
+            "queryName": query_name,
+        }
+
+        return self._post(self.base_path, data=body)
+
+    def get_saved_query_by_id(
+        self,
+        query_id: str,
+    ) -> dict:
+        """Get saved query by query ID.
+
+        Args:
+            query_id (str): The saved query's ID.
+
+        Returns:
+            dict: Saved query result.
+        """
+        return self._get(f"{self.base_path}/{query_id}")
+
+    def edit_saved_query_by_id(
+        self,
+        query_id: str,
+        query: str,
+        query_name: str,
+    ) -> dict:
+        """Edit an existing saved query by query ID.
+
+        Args:
+            query_id (str): The saved query's ID.
+            query (str): New query string.
+            query_name (str): New saved query name.
+
+        Returns:
+            dict: Edited saved query result.
+        """
+        body = {
+            "query": query,
+            "queryName": query_name,
+        }
+
+        return self._put(f"{self.base_path}/{query_id}", data=body)
+
+    def delete_saved_query_by_id(
+        self,
+        query_id: str,
+    ) -> dict:
+        """Delete saved query by query ID.
+
+        Args:
+            query_id (str): The saved query's ID.
+
+        Returns:
+            dict: Delete results.
+        """
+        return self._delete(f"{self.base_path}/{query_id}")

Diff for: censys/cli/commands/account.py

@@ -32,7 +32,7 @@ def cli_account(args: argparse.Namespace):  # pragma: no cover
             quota = account["quota"]
             table.add_row(
                 "Query Quota",
-                f"{quota['used']} / {quota['allowance']} ({quota['used']/quota['allowance'] * 100 :.2f}%)",
+                f"{quota['used']} / {quota['allowance']} ({quota['used']/quota['allowance'] * 100 :.2f}%)",  # noqa
             )
             table.add_row("Quota Resets At", quota["resets_at"])
             console.print(table)

Diff for: censys/cli/commands/asm.py

@@ -11,10 +11,12 @@
 from rich.progress import Progress, TaskID
 from rich.prompt import Confirm, Prompt
 
+from censys.asm.saved_queries import SavedQueries
 from censys.asm.seeds import SEED_TYPES, Seeds
 from censys.cli.utils import console
 from censys.common.config import DEFAULT, get_config, write_config
 from censys.common.exceptions import (
+    CensysAsmException,
     CensysSeedNotFoundException,
     CensysUnauthorizedException,
 )
@@ -437,6 +439,97 @@ def add_seed_arguments(parser: argparse._SubParsersAction, is_delete=False) -> N
     )
 
 
+def cli_list_saved_queries(args: argparse.Namespace):
+    """List saved queries subcommand.
+
+    Args:
+        args (Namespace): Argparse Namespace.
+    """
+    s = SavedQueries(args.api_key)
+    try:
+        res = s.get_saved_queries(
+            args.query_name_prefix, args.page_size, args.page, args.filter_term
+        )
+
+        if args.csv:
+            console.print("queryId,queryName,query,createdAt")
+            for query in res["results"]:
+                console.print(
+                    f"{query['queryId']},{query['queryName']},{query['query']},{query['createdAt']}"  # noqa: E231
+                )
+        else:
+            console.print_json(json.dumps(res))
+    except KeyError:
+        console.print("Failed to get saved queries.")
+        sys.exit(1)
+
+
+def cli_add_saved_query(args: argparse.Namespace):
+    """Add saved query subcommand.
+
+    Args:
+        args (Namespace): Argparse Namespace.
+    """
+    s = SavedQueries(args.api_key)
+    res = s.add_saved_query(args.query, args.query_name)
+    try:
+        console.print(
+            f"Added saved query.\nQuery name: {res['result']['queryName']}\nQuery: {res['result']['query']}\nQuery ID: {res['result']['queryId']}\nCreated at: {res['result']['createdAt']}"
+        )
+    except (KeyError, CensysAsmException):
+        console.print("Failed to add saved query.")
+        sys.exit(1)
+
+
+def cli_get_saved_query_by_id(args: argparse.Namespace):
+    """Get saved query by id subcommand.
+
+    Args:
+        args (Namespace): Argparse Namespace.
+    """
+    s = SavedQueries(args.api_key)
+    res = s.get_saved_query_by_id(args.query_id)
+    try:
+        console.print(
+            f"Query name: {res['result']['queryName']}\nQuery: {res['result']['query']}\nQuery ID: {res['result']['queryId']}\nCreated at: {res['result']['createdAt']}"
+        )
+    except (KeyError, CensysAsmException):
+        console.print("Failed to get saved query.")
+        sys.exit(1)
+
+
+def cli_edit_saved_query_by_id(args: argparse.Namespace):
+    """Edit saved query by id subcommand.
+
+    Args:
+        args (Namespace): Argparse Namespace.
+    """
+    s = SavedQueries(args.api_key)
+    res = s.edit_saved_query_by_id(args.query_id, args.query, args.query_name)
+    try:
+        console.print(
+            f"Edited saved query.\nQuery name: {res['result']['queryName']}\nQuery: {res['result']['query']}\nQuery ID: {res['result']['queryId']}\nCreated at: {res['result']['createdAt']}"
+        )
+    except (KeyError, CensysAsmException):
+        console.print("Failed to edit saved query.")
+        sys.exit(1)
+
+
+def cli_delete_saved_query_by_id(args: argparse.Namespace):
+    """Delete saved query by id subcommand.
+
+    Args:
+        args (Namespace): Argparse Namespace.
+    """
+    s = SavedQueries(args.api_key)
+    res = s.delete_saved_query_by_id(args.query_id)
+    if res == {}:
+        console.print(f"Deleted saved query with ID {args.query_id}.")
+    else:
+        console.print("Failed to delete saved query.")
+        sys.exit(1)
+
+
 def include(parent_parser: argparse._SubParsersAction, parents: dict):
     """Include this subcommand into the parent parser.
 
@@ -568,3 +661,117 @@ def add_verbose(parser):
     )
     add_verbose(list_parser)
     list_parser.set_defaults(func=cli_list_seeds)
+
+    list_saved_queries_parser = asm_subparser.add_parser(
+        "list-saved-queries",
+        description="List all ASM saved queries, optionally filtered by query name prefix and filter term",
+        help="list saved queries",
+        parents=[parents["asm_auth"]],
+    )
+    list_saved_queries_parser.add_argument(
+        "--query-name-prefix",
+        help="Prefix for the saved query name to filter by",
+        type=str,
+        default="",
+    )
+    list_saved_queries_parser.add_argument(
+        "--filter-term",
+        help="Term used to filter the list of saved query names and the saved queries",
+        type=str,
+        default="",
+    )
+    list_saved_queries_parser.add_argument(
+        "--page-size",
+        help="Number of results to return. Defaults to 50.",
+        type=int,
+        default=50,
+    )
+    list_saved_queries_parser.add_argument(
+        "--page",
+        help="Page number to begin at when searching. Defaults to 1.",
+        type=int,
+        default=1,
+    )
+    list_saved_queries_parser.add_argument(
+        "--csv", help="output in CSV format (otherwise JSON)", action="store_true"
+    )
+    add_verbose(list_saved_queries_parser)
+    list_saved_queries_parser.set_defaults(func=cli_list_saved_queries)
+
+    add_saved_query_parser = asm_subparser.add_parser(
+        "add-saved-query",
+        description="Add a saved query to ASM",
+        help="add saved query",
+        parents=[parents["asm_auth"]],
+    )
+    add_saved_query_parser.add_argument(
+        "--query-name",
+        help="Name of the saved query",
+        type=str,
+        required=True,
+    )
+    add_saved_query_parser.add_argument(
+        "--query",
+        help="Query string",
+        type=str,
+        required=True,
+    )
+    add_verbose(add_saved_query_parser)
+    add_saved_query_parser.set_defaults(func=cli_add_saved_query)
+
+    get_saved_query_by_id_parser = asm_subparser.add_parser(
+        "get-saved-query-by-id",
+        description="Get a saved query by ID",
+        help="get saved query by ID",
+        parents=[parents["asm_auth"]],
+    )
+    get_saved_query_by_id_parser.add_argument(
+        "--query-id",
+        help="ID of the saved query",
+        type=str,
+        required=True,
+    )
+    add_verbose(get_saved_query_by_id_parser)
+    get_saved_query_by_id_parser.set_defaults(func=cli_get_saved_query_by_id)
+
+    edit_saved_query_by_id_parser = asm_subparser.add_parser(
+        "edit-saved-query-by-id",
+        description="Edit a saved query by ID",
+        help="edit saved query by ID",
+        parents=[parents["asm_auth"]],
+    )
+    edit_saved_query_by_id_parser.add_argument(
+        "--query-id",
+        help="ID of the saved query",
+        type=str,
+        required=True,
+    )
+    edit_saved_query_by_id_parser.add_argument(
+        "--query-name",
+        help="Name of the saved query",
+        type=str,
+        required=True,
+    )
+    edit_saved_query_by_id_parser.add_argument(
+        "--query",
+        help="Query string",
+        type=str,
+        required=True,
+    )
+    add_verbose(edit_saved_query_by_id_parser)
+    edit_saved_query_by_id_parser.set_defaults(func=cli_edit_saved_query_by_id)
+
+    delete_saved_query_by_id_parser = asm_subparser.add_parser(
+        "delete-saved-query-by-id",
+        description="Delete a saved query by ID",
+        help="delete saved query by ID",
+        parents=[parents["asm_auth"]],
+    )
+    delete_saved_query_by_id_parser.add_argument(
+        "--query-id",
+        help="ID of the saved query",
+        type=str,
+        required=True,
+    )
+    add_verbose(delete_saved_query_by_id_parser)
+    delete_saved_query_by_id_parser.set_defaults(func=cli_delete_saved_query_by_id)