diff --git a/censys/asm/api.py b/censys/asm/api.py index 6e69dae8..6ed614f5 100644 --- a/censys/asm/api.py +++ b/censys/asm/api.py @@ -110,3 +110,11 @@ def _get_logbook_page( args = {"cursor": res["nextCursor"]} yield from res["events"] + + def get_workspace_id(self) -> str: + """Get the workspace ID. + + Returns: + str: The workspace ID. + """ + return self._get("/integrations/v1/account")["workspaceId"] diff --git a/censys/asm/inventory.py b/censys/asm/inventory.py index b046cbb7..e1232423 100644 --- a/censys/asm/inventory.py +++ b/censys/asm/inventory.py @@ -1,4 +1,5 @@ """Interact with the Censys Inventory Search API.""" +import warnings from typing import List, Optional from .api import CensysAsmAPI @@ -11,7 +12,7 @@ class InventorySearch(CensysAsmAPI): def search( self, - workspaces: List[str], + workspaces: Optional[List[str]] = None, query: Optional[str] = None, page_size: Optional[int] = None, cursor: Optional[str] = None, @@ -21,7 +22,7 @@ def search( """Search inventory data. Args: - workspaces (List[str]): List of workspace IDs to search. + workspaces (List[str], optional): List of workspace IDs to search. Deprecated. The workspace associated with `CENSYS-API-KEY` will be used automatically. query (str, optional): Query string. page_size (int, optional): Number of results to return. Defaults to 50. cursor (str, optional): Cursor to start search from. @@ -31,17 +32,31 @@ def search( Returns: dict: Inventory search results. """ + if workspaces is None: + workspaces = [self.get_workspace_id()] + else: + warnings.warn( + "The field 'workspaces' is being deprecated. The workspace associated with `CENSYS-API-KEY` will be used automatically.", + category=DeprecationWarning, + stacklevel=2, + ) if page_size is None: page_size = 50 + args = { "workspaces": workspaces, - "query": query, "pageSize": page_size, - "cursor": cursor, - "sort": sort, - "fields": fields, } + if query: + args["query"] = query + if cursor: + args["cursor"] = cursor + if sort: + args["sort"] = sort + if fields: + args["fields"] = fields + return self._get(self.base_path, args=args) def aggregate( diff --git a/censys/cli/commands/asm.py b/censys/cli/commands/asm.py index a05cd56d..5e019f17 100644 --- a/censys/cli/commands/asm.py +++ b/censys/cli/commands/asm.py @@ -11,6 +11,8 @@ from rich.progress import Progress, TaskID from rich.prompt import Confirm, Prompt +from censys.asm.api import CensysAsmAPI +from censys.asm.inventory import InventorySearch from censys.asm.saved_queries import SavedQueries from censys.asm.seeds import SEED_TYPES, Seeds from censys.cli.utils import console @@ -40,6 +42,14 @@ def cli_asm_config(_: argparse.Namespace): # pragma: no cover api_key = Prompt.ask(api_key_prompt, console=console) or api_key + try: + c = CensysAsmAPI(api_key) + w = c.get_workspace_id() + console.print(f"Successfully authenticated to workspace {w}") + except CensysUnauthorizedException: + console.print("Failed to authenticate") + sys.exit(1) + if not api_key: console.print("Please enter valid credentials") sys.exit(1) @@ -530,6 +540,75 @@ def cli_delete_saved_query_by_id(args: argparse.Namespace): sys.exit(1) +def cli_execute_saved_query_by_name(args: argparse.Namespace): + """Execute saved query by name subcommand. + + Args: + args (Namespace): Argparse Namespace. + """ + s = InventorySearch(args.api_key) + q = SavedQueries(args.api_key) + # get query from name + queries = q.get_saved_queries(args.query_name, 1, 1) + results = queries.get("results") + if not results: + console.print("No saved query found with that name.") + sys.exit(1) + query = results[0]["query"] + + try: + res = s.search(None, query, args.page_size, None, args.sort, args.fields) + console.print_json(json.dumps(res)) + except CensysAsmException: + console.print("Failed to execute saved query.") + sys.exit(1) + + +def cli_execute_saved_query_by_id(args: argparse.Namespace): + """Execute saved query by id subcommand. + + Args: + args (Namespace): Argparse Namespace. + """ + s = InventorySearch(args.api_key) + q = SavedQueries(args.api_key) + try: + query_json = q.get_saved_query_by_id(args.query_id) + query = query_json["result"]["query"] + except (KeyError, CensysAsmException): + console.print("No saved query found with that ID.") + sys.exit(1) + try: + res = s.search(None, query, args.page_size, None, args.sort, args.fields) + console.print_json(json.dumps(res)) + except CensysAsmException: + console.print("Failed to execute saved query.") + sys.exit(1) + + +def cli_search(args: argparse.Namespace): + """Inventory search subcommand. + + Args: + args (Namespace): Argparse Namespace. + """ + s = InventorySearch(args.api_key) + + try: + res = s.search( + args.workspaces, + args.query, + args.page_size, + args.cursor, + args.sort, + args.fields, + ) + console.print_json(json.dumps(res)) + except CensysAsmException: + console.print("Failed to execute query.") + sys.exit(1) + + def include(parent_parser: argparse._SubParsersAction, parents: dict): """Include this subcommand into the parent parser. @@ -775,3 +854,116 @@ def add_verbose(parser): ) add_verbose(delete_saved_query_by_id_parser) delete_saved_query_by_id_parser.set_defaults(func=cli_delete_saved_query_by_id) + + execute_saved_query_by_name_parser = asm_subparser.add_parser( + "execute-saved-query-by-name", + description="Execute a saved query by name in inventory search", + help="execute saved query by name", + parents=[parents["asm_auth"]], + ) + execute_saved_query_by_name_parser.add_argument( + "--query-name", + help="Query name", + type=str, + required=True, + ) + execute_saved_query_by_name_parser.add_argument( + "--page-size", + help="Number of results to return. Defaults to 50.", + type=int, + default=50, + ) + execute_saved_query_by_name_parser.add_argument( + "--sort", + help="Sort order for results", + type=List[str], + default=[], + ) + execute_saved_query_by_name_parser.add_argument( + "--fields", + help="Fields to include in results", + type=List[str], + default=[], + ) + add_verbose(execute_saved_query_by_name_parser) + execute_saved_query_by_name_parser.set_defaults( + func=cli_execute_saved_query_by_name + ) + + execute_saved_query_by_id_parser = asm_subparser.add_parser( + "execute-saved-query-by-id", + description="Execute a saved query by id in inventory search", + help="execute saved query by id", + parents=[parents["asm_auth"]], + ) + execute_saved_query_by_id_parser.add_argument( + "--query-id", + help="Query ID", + type=str, + required=True, + ) + execute_saved_query_by_id_parser.add_argument( + "--page-size", + help="Number of results to return. Defaults to 50.", + type=int, + default=50, + ) + execute_saved_query_by_id_parser.add_argument( + "--sort", + help="Sort order for results", + type=List[str], + default=[], + ) + execute_saved_query_by_id_parser.add_argument( + "--fields", + help="Fields to include in results", + type=List[str], + default=[], + ) + add_verbose(execute_saved_query_by_id_parser) + execute_saved_query_by_id_parser.set_defaults(func=cli_execute_saved_query_by_id) + + search_parser = asm_subparser.add_parser( + "search", + description="Execute a query in inventory search", + help="execute query in inventory search", + parents=[parents["asm_auth"]], + ) + search_parser.add_argument( + "--query", + help="Query string", + type=str, + required=True, + ) + search_parser.add_argument( + "--page-size", + help="Number of results to return. Defaults to 50.", + type=int, + default=50, + ) + search_parser.add_argument( + "--cursor", + help="Cursor to use for pagination", + type=str, + default="", + ) + search_parser.add_argument( + "--sort", + help="Sort order for results", + type=List[str], + default=[], + ) + search_parser.add_argument( + "--fields", + help="Fields to include in results", + type=List[str], + default=[], + ) + search_parser.add_argument( + "--workspaces", + help="Workspace IDs to search. Deprecated. The workspace associated with `CENSYS-API-KEY` will be used automatically.", + type=str, + required=False, + ) + add_verbose(search_parser) + search_parser.set_defaults(func=cli_search) diff --git a/docs/conf.py b/docs/conf.py index a49cb60e..99c849d0 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -63,6 +63,8 @@ # This pattern also affects html_static_path and html_extra_path. exclude_patterns = ["_build", "Thumbs.db", ".DS_Store", "**tests**"] +suppress_warnings = ["autosectionlabel.usage-cli"] + # -- Options for HTML output ------------------------------------------------- diff --git a/docs/requirements.txt b/docs/requirements.txt index 9ce033ad..1ee88c8e 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -1,5 +1,5 @@ sphinx==7.2.6 -sphinxcontrib-autoprogram==0.1.8 +sphinxcontrib-autoprogram==0.1.9 sphinx-rtd-theme==1.3.0 sphinx-prompt==1.8.0 sphinx-tabs==3.4.1 diff --git a/docs/usage-cli.rst b/docs/usage-cli.rst index 6b7b4c57..24e90ae1 100644 --- a/docs/usage-cli.rst +++ b/docs/usage-cli.rst @@ -341,3 +341,36 @@ Below we show an example of deleting a saved query by ID from the CLI. .. prompt:: bash censys asm delete-saved-query-by-id --query-id 'Some ID' + +``execute-saved-query-by-name`` +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +See CLI command :ref:`asm execute-saved-query-by-name` for detail documentation of parameters. + +Below we show an example of executing a saved query by name from the CLI. + +.. prompt:: bash + + censys asm execute-saved-query-by-name --query-name 'Some query name' + +``execute-saved-query-by-id`` +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +See CLI command :ref:`asm execute-saved-query-by-id` for detail documentation of parameters. + +Below we show an example of executing a saved query by ID from the CLI. + +.. prompt:: bash + + censys asm execute-saved-query-by-id --query-id 'Some query ID' + +``search`` +^^^^^^^^^^ + +See CLI command :ref:`asm search` for detail documentation of parameters. + +Below we show an example of executing an inventory search query from the CLI. + +.. prompt:: bash + + censys asm search --query 'Some query' \ No newline at end of file diff --git a/examples/asm/search.py b/examples/asm/search.py new file mode 100644 index 00000000..0c7a2bc0 --- /dev/null +++ b/examples/asm/search.py @@ -0,0 +1,11 @@ +"""Execute a saved query by ID.""" +from censys.asm import InventorySearch + +# Create an instance of the SavedQueries class +inventory_search = InventorySearch() + +# Delete the saved query by ID +query = "host.services.service_name: HTTP" +res = inventory_search.search(query=query) + +print(res) diff --git a/pyproject.toml b/pyproject.toml index 74073c17..77d94abe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "censys" -version = "2.2.12" +version = "2.2.13" description = "An easy-to-use and lightweight API wrapper for Censys APIs (censys.io)." authors = ["Censys, Inc. "] license = "Apache-2.0" diff --git a/tests/asm/test_api.py b/tests/asm/test_api.py index 614a930a..1459e20b 100644 --- a/tests/asm/test_api.py +++ b/tests/asm/test_api.py @@ -108,3 +108,17 @@ def keyword_callback(request): res = list(self.api._get_page(f"/{keyword}")) # Assertions assert res == page_json[keyword] + second_page + + def test_get_workspace_id(self): + self.responses.add( + responses.GET, + f"{self.base_url}/integrations/v1/account", + status=200, + json={"workspaceId": "test-workspace-id"}, + ) + + # Actual call + res = self.api.get_workspace_id() + + # Assertions + assert res == "test-workspace-id" diff --git a/tests/asm/test_inventory.py b/tests/asm/test_inventory.py index 4baddd6d..be25fb33 100644 --- a/tests/asm/test_inventory.py +++ b/tests/asm/test_inventory.py @@ -2,7 +2,7 @@ from parameterized import parameterized from ..utils import CensysTestCase -from .utils import BASE_URL +from .utils import BASE_URL, WORKSPACE_ID from censys.asm.inventory import InventorySearch INVENTORY_BASE_PATH = f"{BASE_URL}/inventory/v1" @@ -47,14 +47,13 @@ def setUp(self): [ ( { - "workspaces": ["1", "2"], "query": "test", "page_size": 50, "cursor": "test", "sort": ["test"], "fields": ["test"], }, - "?workspaces=1&workspaces=2&query=test&pageSize=50&cursor=test&sort=test&fields=test", + "?workspaces=test-workspace-id&query=test&pageSize=50&cursor=test&sort=test&fields=test", ), ( { @@ -66,6 +65,9 @@ def setUp(self): ] ) def test_search(self, kwargs, params): + mock_request = self.mocker.patch("censys.asm.api.CensysAsmAPI.get_workspace_id") + mock_request.return_value = WORKSPACE_ID + # Setup response self.responses.add( responses.GET, diff --git a/tests/asm/utils.py b/tests/asm/utils.py index 73542db9..08d90967 100644 --- a/tests/asm/utils.py +++ b/tests/asm/utils.py @@ -6,6 +6,7 @@ V2_URL = f"{BASE_URL}/v2" BETA_URL = f"{BASE_URL}/beta" INVENTORY_URL = f"{BASE_URL}/inventory" +WORKSPACE_ID = "test-workspace-id" class MockResponse: diff --git a/tests/cli/test_asm.py b/tests/cli/test_asm.py index 8cadb682..95701e4c 100644 --- a/tests/cli/test_asm.py +++ b/tests/cli/test_asm.py @@ -8,7 +8,7 @@ from responses import matchers from responses.matchers import json_params_matcher -from tests.asm.utils import INVENTORY_URL, V1_URL +from tests.asm.utils import INVENTORY_URL, V1_URL, WORKSPACE_ID from tests.cli.test_config import TEST_CONFIG_PATH from tests.utils import CensysTestCase @@ -145,6 +145,27 @@ "totalResults": 1, } +SEARCH_JSON = { + "totalHits": 3, + "queryDurationMillis": 50, + "hits": [ + {"_details": {}, "domain": {"name": "foo.com"}, "type": "DOMAIN"}, + { + "_details": {}, + "domain": {"name": "bar.foo.com"}, + "type": "DOMAIN", + }, + { + "_details": {}, + "domain": {"name": "b.foo.com"}, + "type": "DOMAIN", + }, + ], +} + +SAVED_QUERY_ID = "12345" +SAVED_QUERY_NAME = "Test query" + TEST_XML_PATH = os.path.join(os.path.dirname(__file__), "test.xml") @@ -1854,3 +1875,319 @@ def test_delete_saved_query_by_id_failed(self): # Actual call with pytest.raises(SystemExit, match="1"): cli_main() + + def test_execute_saved_query_by_name(self): + # Mock + mock_request = self.mocker.patch("censys.asm.api.CensysAsmAPI.get_workspace_id") + mock_request.return_value = WORKSPACE_ID + mock_query = self.mocker.patch("censys.asm.SavedQueries.get_saved_queries") + mock_query.return_value = { + "results": [ + { + "queryId": "1", + "queryName": "foo domain", + "query": "domain: foo.com", + "createdAt": "2024-01-01T01:00:00.000Z", + } + ] + } + self.patch_args( + [ + "censys", + "asm", + "execute-saved-query-by-name", + "--query-name", + "foo domain", + ], + asm_auth=True, + ) + self.responses.add( + responses.GET, + INVENTORY_URL + "/v1", + status=200, + json=SEARCH_JSON, + match=[ + matchers.query_param_matcher( + { + "workspaces": WORKSPACE_ID, + "query": "domain: foo.com", + "pageSize": 50, + } + ) + ], + ) + + # Actual call + temp_stdout = StringIO() + with contextlib.redirect_stdout(temp_stdout): + cli_main() + + # Assertions + actual_json = json.loads(temp_stdout.getvalue()) + assert actual_json == SEARCH_JSON + + def test_execute_saved_query_by_name_not_found(self): + # Mock + mock_query = self.mocker.patch("censys.asm.SavedQueries.get_saved_queries") + mock_query.return_value = {"results": []} + self.patch_args( + [ + "censys", + "asm", + "execute-saved-query-by-name", + "--query-name", + "nonexistent query name", + ], + asm_auth=True, + ) + + # Actual call + temp_stdout = StringIO() + with pytest.raises(SystemExit, match="1"), contextlib.redirect_stdout( + temp_stdout + ): + cli_main() + + # Assertions + assert "No saved query found with that name." in temp_stdout.getvalue() + + def test_execute_saved_query_by_name_failed(self): + # Mock + mock_request = self.mocker.patch("censys.asm.api.CensysAsmAPI.get_workspace_id") + mock_request.return_value = WORKSPACE_ID + mock_query = self.mocker.patch("censys.asm.SavedQueries.get_saved_queries") + mock_query.return_value = { + "results": [ + { + "queryId": "1", + "queryName": "foo domain", + "query": "domain: foo.com", + "createdAt": "2024-01-01T01:00:00.000Z", + } + ] + } + self.patch_args( + [ + "censys", + "asm", + "execute-saved-query-by-name", + "--query-name", + "foo domain", + ], + asm_auth=True, + ) + self.responses.add( + responses.GET, + INVENTORY_URL + "/v1", + status=500, + json={}, + match=[ + matchers.query_param_matcher( + { + "workspaces": WORKSPACE_ID, + "query": "domain: foo.com", + "pageSize": 50, + } + ) + ], + ) + + # Actual call + temp_stdout = StringIO() + with pytest.raises(SystemExit, match="1"), contextlib.redirect_stdout( + temp_stdout + ): + cli_main() + + # Assertions + assert "Failed to execute saved query." in temp_stdout.getvalue() + + def test_execute_saved_query_by_id(self): + # Mock + mock_request = self.mocker.patch("censys.asm.api.CensysAsmAPI.get_workspace_id") + mock_request.return_value = WORKSPACE_ID + mock_query = self.mocker.patch("censys.asm.SavedQueries.get_saved_query_by_id") + mock_query.return_value = SAVED_QUERY_JSON + self.patch_args( + [ + "censys", + "asm", + "execute-saved-query-by-id", + "--query-id", + "100", + ], + asm_auth=True, + ) + self.responses.add( + responses.GET, + INVENTORY_URL + "/v1", + status=200, + json=SEARCH_JSON, + match=[ + matchers.query_param_matcher( + { + "workspaces": WORKSPACE_ID, + "query": "domain: foo.com", + "pageSize": 50, + } + ) + ], + ) + + # Actual call + temp_stdout = StringIO() + with contextlib.redirect_stdout(temp_stdout): + cli_main() + + # Assertions + actual_json = json.loads(temp_stdout.getvalue()) + assert actual_json == SEARCH_JSON + + def test_execute_saved_query_by_id_not_found(self): + # Mock + mock_query = self.mocker.patch("censys.asm.SavedQueries.get_saved_query_by_id") + mock_query.return_value = { + "status_code": 404, + "error": "No query found with ID: 500", + } + self.patch_args( + [ + "censys", + "asm", + "execute-saved-query-by-id", + "--query-id", + "500", + ], + asm_auth=True, + ) + + # Actual call + temp_stdout = StringIO() + with pytest.raises(SystemExit, match="1"), contextlib.redirect_stdout( + temp_stdout + ): + cli_main() + + # Assertions + assert "No saved query found with that ID." in temp_stdout.getvalue() + + def test_execute_saved_query_by_id_failed(self): + # Mock + mock_request = self.mocker.patch("censys.asm.api.CensysAsmAPI.get_workspace_id") + mock_request.return_value = WORKSPACE_ID + mock_query = self.mocker.patch("censys.asm.SavedQueries.get_saved_query_by_id") + mock_query.return_value = SAVED_QUERY_JSON + self.patch_args( + [ + "censys", + "asm", + "execute-saved-query-by-id", + "--query-id", + "100", + ], + asm_auth=True, + ) + self.responses.add( + responses.GET, + INVENTORY_URL + "/v1", + status=500, + json={}, + match=[ + matchers.query_param_matcher( + { + "workspaces": WORKSPACE_ID, + "query": "domain: foo.com", + "pageSize": 50, + } + ) + ], + ) + + # Actual call + temp_stdout = StringIO() + with pytest.raises(SystemExit, match="1"), contextlib.redirect_stdout( + temp_stdout + ): + cli_main() + + # Assertions + assert "Failed to execute saved query." in temp_stdout.getvalue() + + def test_search(self): + # Mock + mock_request = self.mocker.patch("censys.asm.api.CensysAsmAPI.get_workspace_id") + mock_request.return_value = WORKSPACE_ID + self.patch_args( + [ + "censys", + "asm", + "search", + "--query", + "domain: foo.com", + ], + asm_auth=True, + ) + self.responses.add( + responses.GET, + INVENTORY_URL + "/v1", + status=200, + json=SEARCH_JSON, + match=[ + matchers.query_param_matcher( + { + "workspaces": WORKSPACE_ID, + "query": "domain: foo.com", + "pageSize": 50, + } + ) + ], + ) + + # Actual call + temp_stdout = StringIO() + with contextlib.redirect_stdout(temp_stdout): + cli_main() + + # Assertions + actual_json = json.loads(temp_stdout.getvalue()) + assert actual_json == SEARCH_JSON + + def test_search_failed(self): + # Mock + mock_request = self.mocker.patch("censys.asm.api.CensysAsmAPI.get_workspace_id") + mock_request.return_value = WORKSPACE_ID + self.patch_args( + [ + "censys", + "asm", + "search", + "--query", + "domain: foo.com", + ], + asm_auth=True, + ) + self.responses.add( + responses.GET, + INVENTORY_URL + "/v1", + status=500, + json={}, + match=[ + matchers.query_param_matcher( + { + "workspaces": WORKSPACE_ID, + "query": "domain: foo.com", + "pageSize": 50, + } + ) + ], + ) + + # Actual call + temp_stdout = StringIO() + with pytest.raises(SystemExit, match="1"), contextlib.redirect_stdout( + temp_stdout + ): + cli_main() + + # Assertions + assert "Failed to execute query." in temp_stdout.getvalue()