feat(cli): Added virtual hosts argument to search

Author: thehappydinoa

censys/cli/commands/search.py

@@ -101,6 +101,9 @@ def cli_search(args: argparse.Namespace):
         if args.pages:
             search_args["pages"] = args.pages
 
+        if args.virtual_hosts:
+            search_args["virtual_hosts"] = args.virtual_hosts
+
         with err_console.status("Searching"):
             query = index.search(args.query, **search_args)
 
@@ -183,6 +186,14 @@ def include(parent_parser: argparse._SubParsersAction, parents: dict):
         type=int,
         help="number of pages of results to return (when set to -1 returns all pages available)",
     )
+    v2_group.add_argument(
+        "--virtual-hosts",
+        type=str,
+        default="EXCLUDE",
+        choices=["INCLUDE", "EXCLUDE", "ONLY"],
+        metavar="INCLUDE|EXCLUDE|ONLY",
+        help="whether to include virtual hosts in the results",
+    )
 
     v1_group = search_parser.add_argument_group(
         f"v1 specific arguments ({', '.join(V1_INDEXES)})"

tests/cli/test_search.py

@@ -341,7 +341,42 @@ def test_max_records(self):
     def test_write_screen_v2(self):
         self.responses.add(
             responses.GET,
-            V2_URL + "/hosts/search?q=service.service_name: HTTP&per_page=100",
+            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()
+
+        json_response = json.loads(temp_stdout.getvalue().strip())
+
+        assert json_response == SEARCH_HOSTS_JSON["result"]["hits"]
+
+    @patch(
+        "argparse._sys.argv",
+        [
+            "censys",
+            "search",
+            "service.service_name: HTTP",
+            "--index-type",
+            "hosts",
+            "--format",
+            "screen",
+            "--pages",
+            "1",
+            "--virtual-hosts",
+            "ONLY",
+        ]
+        + CensysTestCase.cli_args,
+    )
+    def test_search_virtual_hosts(self):
+        self.responses.add(
+            responses.GET,
+            V2_URL
+            + "/hosts/search?q=service.service_name%3A+HTTP&per_page=100&virtual_hosts=ONLY",
             status=200,
             json=SEARCH_HOSTS_JSON,
         )