Add JSON output support (#150)

* Output json

* Tidy

* Remove strenum

Author: inverse

cert_host_scraper/cli.py

@@ -1,4 +1,5 @@
 import asyncio
+import json
 import logging
 import sys
 
@@ -31,6 +32,15 @@ def validate_status_code(
         return NO_STATUS_CODE_FILTER
 
 
+class Output:
+    TABLE = "table"
+    JSON = "json"
+
+    @classmethod
+    def values(cls) -> list:
+        return [cls.TABLE, cls.JSON]
+
+
 @click.group()
 @click.option("--debug", is_flag=True, help="Whether to enable debug level output")
 @click.version_option(__version__, message="%(version)s")
@@ -61,21 +71,28 @@ def cli(debug: bool):
     help="Number of URLs to process at once",
     default=20,
 )
+@click.option(
+    "--output", type=click.Choice(Output.values()), required=True, default="table"
+)
 def search(
     search: str,
     status_code: int,
     timeout: int,
     clean: bool,
     strip: bool,
     batch_size: int,
+    output: str,
 ):
     """
     Search the certificate transparency log.
     """
     if strip:
         search = strip_url(search)
 
-    click.echo(f"Searching for {search}")
+    display_json = output == Output.JSON
+
+    if not display_json:
+        click.echo(f"Searching for {search}")
     options = Options(timeout, clean)
     results = []
     try:
@@ -84,11 +101,12 @@ def search(
         click.echo(f"Failed to search for results: {e}")
         sys.exit(1)
 
-    click.echo(f"Found {len(urls)} URLs for {search}")
+    if not display_json:
+        click.echo(f"Found {len(urls)} URLs for {search}")
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
     chunks = list(divide_chunks(urls, batch_size))
-    for chunk_index in track(range(len(chunks)), "Checking URLs"):
+    for chunk_index in track(range(len(chunks)), "Checking URLs", disable=display_json):
         chunk_result = loop.run_until_complete(
             asyncio.gather(*[validate_url(url, options) for url in chunks[chunk_index]])
         )
@@ -100,23 +118,30 @@ def search(
     else:
         display = result.scraped
 
-    table = Table(show_header=True, header_style="bold", box=box.MINIMAL)
-    table.add_column("URL")
-    table.add_column("Status Code")
-    for url_result in display:
-        display_code = str(url_result.status_code)
-        if url_result.status_code == -1:
-            display_code = "-"
-
-        url = url_result.url
-        if url_result.status_code == 200:
-            display_code = f"[green]{display_code}[/green]"
-            url = f"[green]{url}[/green]"
-
-        table.add_row(url, display_code)
-
-    console = Console()
-    console.print(table)
+    if display_json:
+        json_output = [
+            {"url": url_result.url, "status_code": url_result.status_code}
+            for url_result in display
+        ]
+        click.echo(json.dumps(json_output, indent=2))
+    else:
+        table = Table(show_header=True, header_style="bold", box=box.MINIMAL)
+        table.add_column("URL")
+        table.add_column("Status Code")
+        for url_result in display:
+            display_code = str(url_result.status_code)
+            if url_result.status_code == -1:
+                display_code = "-"
+
+            url = url_result.url
+            if url_result.status_code == 200:
+                display_code = f"[green]{display_code}[/green]"
+                url = f"[green]{url}[/green]"
+
+            table.add_row(url, display_code)
+
+        console = Console()
+        console.print(table)
 
 
 if __name__ == "__main__":

tests/test_cli.py

@@ -43,3 +43,8 @@ def test_search_upper_invalid_status_code(self):
         runner = CliRunner()
         result = runner.invoke(cli, ["search", "example.com", "--status-code", "600"])
         self.assertEqual(result.exit_code, 2)
+
+    def test_invalid_output(self):
+        runner = CliRunner()
+        result = runner.invoke(cli, ["search", "example.com", "--output", "csv"])
+        self.assertEqual(result.exit_code, 2)