Merge pull request #266 from censys/grace-murphy/censys-search-csv-fix

Edit output format options for Search/View CLI

Author: grace-murphy

censys/cli/commands/search.py

@@ -71,6 +71,14 @@ def cli_search(args: argparse.Namespace):
         if args.max_records:
             search_args["max_records"] = args.max_records
 
+        if args.output:
+            if args.output.endswith(".csv"):
+                write_args["file_format"] = "csv"
+            else:
+                write_args["file_format"] = "json"
+        else:
+            write_args["file_format"] = "screen"
+
         fields: List[str] = []
         if args.fields:
             if args.overwrite:
@@ -91,12 +99,17 @@ def cli_search(args: argparse.Namespace):
         with err_console.status("Searching"):
             results = list(index.search(args.query, **search_args))
     elif index_type in V2_INDEXES:
-        if args.format == "csv":
+        if args.format == "csv" or (args.output and not args.output.endswith(".json")):
             raise CensysCLIException(
-                f"CSV output is not supported for the {index_type} index."
+                "JSON is the only valid file format for Search 2.0 responses."
             )
         index = getattr(c.v2, index_type)
 
+        if args.output:
+            write_args["file_format"] = "json"
+        else:
+            write_args["file_format"] = "screen"
+
         if args.pages:
             search_args["pages"] = args.pages
 
@@ -162,7 +175,7 @@ def include(parent_parser: argparse._SubParsersAction, parents: dict):
         default="screen",
         choices=["screen", "json", "csv"],
         metavar="screen|json|csv",
-        help="format of output (csv is only supported for the certificates index)",
+        help=argparse.SUPPRESS,
     )
     search_parser.add_argument(
         "-o",

censys/cli/commands/view.py

@@ -36,7 +36,6 @@ def cli_view(args: argparse.Namespace):
     write_args = {
         "file_format": "json" if args.output else "screen",
         "file_path": args.output,
-        "base_name": f"censys-view-{args.document_id}",
     }
 
     if args.at_time:

censys/cli/utils.py

@@ -86,9 +86,8 @@ def _write_screen(search_results: Results):  # pragma: no cover
 
 def write_file(
     results_list: Results,
-    file_format: str = "screen",
+    file_format: Optional[str] = None,
     file_path: Optional[str] = None,
-    base_name: str = "censys-query-output",
     csv_fields: List[str] = [],
 ):
     """Maps formats and writes results.
@@ -97,17 +96,13 @@ def write_file(
         results_list (Results): A list of results from the API query.
         file_format (str): Optional; The format of the output.
         file_path (str): Optional; A path to write results to.
-        base_name (str): Optional; The base name of the output file.
         csv_fields (List[str]): Optional; A list of fields to write to CSV.
     """
     if file_format and isinstance(file_format, str):
         file_format = file_format.lower()
 
     if not file_path:
-        # This method just creates some dynamic file names
-        file_path = ".".join([base_name, file_format])
-    elif file_path.endswith(".json"):
-        file_format = "json"
+        file_path = "temp-out.json"
 
     if file_format == "json":
         _write_json(file_path, results_list)

docs/usage-cli.rst

@@ -44,11 +44,11 @@ Below we show an example of Viewing a host from the CLI.
 
     censys view 8.8.8.8
 
-You can save results to a file using the ``-f`` and ``-o`` arguments.
+You can save results to a file using the ``-o`` argument.
 
 .. prompt:: bash
 
-    censys view 8.8.8.8 -f json -o google.json
+    censys view 8.8.8.8 -o google.json
 
 We can then parse this json with something like ``jq``.
 

tests/cli/test_search.py

@@ -55,8 +55,8 @@ def test_no_creds(self, mock_file):
             "certs",
             "--fields",
             "parsed.issuer.country",
-            "--format",
-            "json",
+            "--output",
+            "censys-certs.json",
         ]
         + CensysTestCase.cli_args,
     )
@@ -77,7 +77,7 @@ def test_write_json(self):
 
         json_path = cli_response.replace(WROTE_PREFIX, "").strip()
         assert json_path.endswith(".json")
-        assert "censys-query-output." in json_path
+        assert "censys-certs." in json_path
 
         with open(json_path) as json_file:
             json_response = json.load(json_file)
@@ -98,8 +98,8 @@ def test_write_json(self):
             "certs",
             "--fields",
             "protocols",
-            "--format",
-            "csv",
+            "--output",
+            "censys-certs.csv",
         ]
         + CensysTestCase.cli_args,
     )
@@ -120,7 +120,7 @@ def test_write_csv(self):
 
         csv_path = cli_response.replace(WROTE_PREFIX, "").strip()
         assert csv_path.endswith(".csv")
-        assert "censys-query-output." in csv_path
+        assert "censys-certs." in csv_path
 
         with open(csv_path) as csv_file:
             csv_reader = csv.reader(csv_file)
@@ -141,35 +141,38 @@ def test_write_csv(self):
             "certs",
             "--fields",
             "parsed.issuer.country",
-            "--format",
-            "json",
             "--output",
-            "censys-certs.json",
+            "censys-certs.html",
         ]
         + CensysTestCase.cli_args,
     )
-    def test_write_output_path(self):
+    def test_write_invalid_output_path(self):
         self.responses.add_callback(
             responses.POST,
             V1_URL + "/search/certificates",
             callback=search_callback,
             content_type="application/json",
         )
 
-        output_path = "censys-certs.json"
+        temp_stdout = StringIO()
+        with contextlib.redirect_stdout(temp_stdout):
+            cli_main()
 
-        cli_main()
+        cli_response = temp_stdout.getvalue().strip()
+        assert cli_response.startswith(WROTE_PREFIX)
 
-        assert os.path.isfile(output_path)
+        json_path = cli_response.replace(WROTE_PREFIX, "").strip()
+        assert json_path.endswith(".html")
+        assert "censys-certs." in json_path
 
-        with open(output_path) as json_file:
+        with open(json_path) as json_file:
             json_response = json.load(json_file)
 
         assert len(json_response) >= 1
         assert "parsed.issuer.country" in json_response[0]
 
         # Cleanup
-        os.remove(output_path)
+        os.remove(json_path)
 
     @patch(
         "argparse._sys.argv",
@@ -181,8 +184,6 @@ def test_write_output_path(self):
             "certs",
             "--fields",
             "443.https.get.headers.server",
-            "--format",
-            "screen",
         ]
         + CensysTestCase.cli_args,
     )
@@ -220,8 +221,6 @@ def test_write_screen(self):
             "443.https.tls.cipher_suite.name",
             "443.https.get.title",
             "443.https.get.headers.server",
-            "--format",
-            "screen",
         ]
         + CensysTestCase.cli_args,
     )
@@ -300,8 +299,6 @@ def test_field_max(self):
             "parsed.names: censys.io",
             "--index-type",
             "certs",
-            "--format",
-            "screen",
             "--max-records",
             "2",
         ]
@@ -331,8 +328,6 @@ def test_max_records(self):
             "service.service_name: HTTP",
             "--index-type",
             "hosts",
-            "--format",
-            "screen",
             "--pages",
             "1",
         ]
@@ -363,8 +358,6 @@ def test_write_screen_v2(self):
             "service.service_name: HTTP",
             "--index-type",
             "hosts",
-            "--format",
-            "screen",
             "--pages",
             "1",
             "--virtual-hosts",
@@ -397,17 +390,61 @@ def test_search_virtual_hosts(self):
             "service.service_name: HTTP",
             "--index-type",
             "hosts",
-            "--format",
-            "csv",
             "--pages",
             "1",
+            "--output",
+            "censys-hosts.json",
+        ]
+        + CensysTestCase.cli_args,
+    )
+    def test_write_json_v2(self):
+        self.responses.add(
+            responses.GET,
+            V2_URL
+            + "/hosts/search?q=service.service_name: HTTP&per_page=100&virtual_hosts=EXCLUDE",
+            status=200,
+            json=SEARCH_HOSTS_JSON,
+        )
+
+        temp_stdout = StringIO()
+        with contextlib.redirect_stdout(temp_stdout):
+            cli_main()
+
+        cli_response = temp_stdout.getvalue().strip()
+        assert cli_response.startswith(WROTE_PREFIX)
+
+        json_path = cli_response.replace(WROTE_PREFIX, "").strip()
+        assert json_path.endswith(".json")
+        assert "censys-hosts." in json_path
+
+        with open(json_path) as json_file:
+            json_response = json.load(json_file)
+
+        assert len(json_response) >= 1
+        assert json_response == SEARCH_HOSTS_JSON["result"]["hits"]
+
+        # Cleanup
+        os.remove(json_path)
+
+    @patch(
+        "argparse._sys.argv",
+        [
+            "censys",
+            "search",
+            "service.service_name: HTTP",
+            "--index-type",
+            "hosts",
+            "--pages",
+            "1",
+            "--output",
+            "censys-hosts.csv",
         ]
         + CensysTestCase.cli_args,
     )
     def test_write_csv_v2(self):
         with pytest.raises(
             CensysCLIException,
-            match="CSV output is not supported for the hosts index.",
+            match="JSON is the only valid file format for Search 2.0 responses.",
         ):
             cli_main()