Merge pull request #11 from Veridise/nikos/async-artifact-improvements

Add script to wait and download an artifact given its step-code and name.

Author: NikosChondros

audithub_client/__main__.py

@@ -12,6 +12,7 @@
     create_version_via_local_archive,
 )
 from .scripts.create_version_via_url import create_version_via_url  # noqa
+from .scripts.download_artifact import download_artifact  # noqa
 from .scripts.get_configuration import get_configuration  # noqa
 from .scripts.get_my_organizations import get_my_organizations  # noqa
 from .scripts.get_my_profile import get_my_profile  # noqa
@@ -31,6 +32,18 @@
 )
 
 
+class LevelFormatter(logging.Formatter):
+    def __init__(self, formats, default_fmt=None):
+        super().__init__()
+        self.formats = formats
+        self.default_fmt = default_fmt or "%(levelname)s: %(message)s"
+
+    def format(self, record):
+        fmt = self.formats.get(record.levelno, self.default_fmt)
+        formatter = logging.Formatter(fmt)
+        return formatter.format(record)
+
+
 @app.meta.default
 def meta(
     *tokens: Annotated[str, Parameter(show=False, allow_leading_hyphen=True)],
@@ -66,15 +79,31 @@ def meta(
         ),
     ] = "INFO",
 ):
+    # Leave all levels except DEBUG "simple", without outputting the module name.
+    # At debug level, also show the module name that produces the log message.
+    default_log_format = "%(asctime)s %(levelname)s %(message)s"
+    debug_log_format = "%(asctime)s %(levelname)s %(name)s %(message)s"
+
+    # Setup logging as usual
     logging.basicConfig(
         level=log_level,
-        # format="%(asctime)s.%(msecs)03d %(filename)s%(name)s %(levelname)s %(message)s",  # cspell:disable-line
-        format="%(asctime)s %(levelname)s %(message)s",  # cspell:disable-line
+        format=default_log_format,
         datefmt="%H:%M:%S",
         stream=sys.stderr,
     )
-    httpx_logger = logging.getLogger("httpx")
-    httpx_logger.setLevel(logging.WARNING)
+    # Now modify the root logger
+    root_logger = logging.getLogger()
+    handler = root_logger.handlers[0]  # basicConfig creates one handler by default
+    handler.setFormatter(
+        LevelFormatter(
+            {logging.DEBUG: debug_log_format}, default_fmt=default_log_format
+        )
+    )
+
+    # For these modules, raise the log level to reduce noise
+    for module in ["httpx", "httpcore"]:
+        module_logger = logging.getLogger(module)
+        module_logger.setLevel(logging.WARNING)
 
     command, bound, _ignored = app.parse_args(tokens)
     # When this script runs with no args, help_print is automatically invoked

audithub_client/api/get_version_comments.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 from dataclasses import dataclass
-from datetime import datetime, timedelta, timezone
+from datetime import datetime
 from typing import Optional
 
 from audithub_client.library.utils import get_dict_of_fields_except

audithub_client/library/auth.py

@@ -32,15 +32,15 @@ def get_access_token(
         "scope": "openid profile",
         "grant_type": "client_credentials",
     }
-    logger.debug("Payload is %s", payload)
+    # logger.debug("Payload is %s", payload)
     response = post(token_url, data=payload, timeout=AUTHENTICATION_TIMEOUT)
     if response.status_code != 200:
         raise RuntimeError(
             f'Failed to get token for client {payload["client_id"]} status = {response.status_code} response ={response.text}'
         )
     token_data = response.json()
     end_time = time.perf_counter()
-    logger.debug(json.dumps(token_data, indent=4))
+    logger.debug(json.dumps(token_data))
     if token_time_listener:
         token_time_listener(end_time - begin_time)
     return token_data["access_token"]

audithub_client/library/http_download.py

@@ -0,0 +1,17 @@
+from pathlib import Path
+from typing import Tuple
+
+from httpx import Client
+
+from audithub_client.library.net_utils import download_file
+
+from .auth import DEFAULT_REQUEST_TIMEOUT
+
+
+def download_from_url(
+    url: str, output_file: Path, timeout=DEFAULT_REQUEST_TIMEOUT
+) -> Tuple[int, str]:
+    """This is a `curl -o`/`wget -O` equivalent"""
+    with Client(timeout=timeout) as client:
+        with client.stream("GET", url) as response:
+            return download_file(response, output_file)

audithub_client/library/invocation_common.py

@@ -60,3 +60,5 @@
         help="An optional task name for this task. If not specified, one will automatically be generated by AuditHub."
     ),
 ]
+
+BooleanArg = Annotated[bool, Parameter(negative_bool=())]

audithub_client/library/json_dump.py

@@ -6,17 +6,18 @@
 from tabulate import tabulate
 
 OutputType = Annotated[
-    Literal["raw", "json", "json-pretty", "pprint", "list", "table"],
+    Literal["raw", "json", "json-pretty", "pprint", "list", "table", "none"],
     Parameter(
         help="""\
-The output format. Options are: 'raw': Python print(), 'json': single-line JSON, 'json-pretty': multi-line JSON, 'pprint': Python pprint(), 'list': list element per line, 'table': tabular view.
+The output format. Options are: 'raw': Python print(), 'json': single-line JSON, 'json-pretty': multi-line JSON, 'pprint': Python pprint(), 'list': list element per line, 'table': tabular view, 'none': omit output completely.
 """
     ),
 ]
 
 
 def dump_dict(document: Any, section: str | None = None, output: OutputType = "json"):
-
+    if output == "none":
+        return
     if section == "sections":
         document = sorted(document.keys())
     elif section:

audithub_client/scripts/download_artifact.py

@@ -0,0 +1,97 @@
+import logging
+import sys
+from pathlib import Path
+from time import sleep
+from typing import Annotated
+
+from cyclopts import Parameter, validators
+
+from ..api.get_task_info import GetTaskInfoArgs, api_get_task_info
+from ..library.http_download import download_from_url
+from ..library.invocation_common import (
+    AuditHubContextType,
+    OrganizationIdType,
+    TaskIdType,
+    app,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@app.command
+def download_artifact(
+    *,
+    organization_id: OrganizationIdType,
+    task_id: TaskIdType,
+    step_code: str,
+    name: str,
+    output_file: Path,
+    timeout: Annotated[int, Parameter(validator=validators.Number(gt=0))] = 30,
+    rpc_context: AuditHubContextType,
+):
+    """
+    Download an artifact by name, potentially waiting for it to become available.
+
+    This is an improved version of get-task-artifact, that potentially waits for the artifact to become
+    available and downloads it using the provided URL.
+
+    Note that you can use 'ah get-task-info' to obtain the list of step_codes
+    (key 'steps') and produced artifacts (key 'artifacts').
+
+    Also note that, due to the asynchronous nature of AuditHub, artifacts may take a short amount of time
+    until they become available, even when the task has finished.
+    This is normal, and this command takes this into account.
+
+    Parameters
+    ----------
+    step_code:
+        The code of the workflow step that produced the artifact
+    name:
+        The name of the artifact.
+    output_file:
+        The local file name to store the output in.
+    timeout:
+        The number of seconds to potentially wait for the artifact to become available.
+    """
+    try:
+        found = False
+        for attempt in range(1, timeout + 1):
+            logger.debug("Starting attempt %d", attempt)
+            task_info = api_get_task_info(
+                rpc_context,
+                GetTaskInfoArgs(organization_id=organization_id, task_id=task_id),
+            )
+            matched_artifacts = [
+                e
+                for e in task_info.get("artifacts", list())
+                if e.get("step_code") == step_code and e.get("name") == name
+            ]
+            if len(matched_artifacts) > 1:
+                # This should be impossible
+                raise RuntimeError(
+                    f"Multiple artifacts matched the condition, bailing: '{matched_artifacts}'"
+                )
+            if len(matched_artifacts) == 1:
+                logger.debug("Artifact found at attempt %d, downloading...", attempt)
+                bytes_written, hr_size = download_from_url(
+                    matched_artifacts[0]["presigned_url"], output_file
+                )
+                logger.info(
+                    f"Downloaded {bytes_written} bytes ({hr_size}) as {output_file}."
+                )
+                found = True
+                break
+            else:
+                logger.info(
+                    "Artifact not found, waiting a sec at attempt %d..", attempt
+                )
+                sleep(1)
+
+        if found:
+            logger.debug("Finished.")
+        else:
+            logger.error("Artifact not found (yet?).")
+            sys.exit(1)
+
+    except Exception as ex:
+        logger.error("Error %s", str(ex), exc_info=ex)

audithub_client/scripts/get_task_artifact.py

@@ -23,7 +23,7 @@ def get_task_artifact(
     rpc_context: AuditHubContextType,
 ):
     """
-    DownGet logs of a task's step.
+    Download a task's artifact, by its id.
 
     Parameters
     ----------

audithub_client/scripts/get_task_info.py

@@ -5,6 +5,7 @@
 from ..api.get_task_info import GetTaskInfoArgs, api_get_task_info
 from ..library.invocation_common import (
     AuditHubContextType,
+    BooleanArg,
     OrganizationIdType,
     TaskIdType,
     app,
@@ -21,7 +22,8 @@ def get_task_info(
     organization_id: OrganizationIdType,
     task_id: TaskIdType,
     output: OutputType = "json",
-    verify: bool = False,
+    verify: BooleanArg = False,
+    check_completed: BooleanArg = False,
     rpc_context: AuditHubContextType,
 ):
     """
@@ -37,6 +39,9 @@ def get_task_info(
         If true, the findings counters are summed. If the sum if zero, the exit code is 0, otherwise it is 1.
         i.e., an exit code of 1 means there is at least one finding in one of the findings categories.
         This argument is independent of any output arguments.
+    check_completed:
+        If true, all steps that are "tools" are checked to see if they reported completed == false.
+        In such a case, we produce a 'timeout error' with an exit code of 2
     """
     try:
         rpc_input = GetTaskInfoArgs(organization_id=organization_id, task_id=task_id)
@@ -60,6 +65,18 @@ def get_task_info(
                         exit_code = 1
             if exit_code == 0:
                 print("No findings reported by this task, exiting with 0")
+        if exit_code == 0 and check_completed:
+            for step in task_info.get("steps", list()):
+                completed = step.get("completed_without_timeout", None)
+                if completed is not None:
+                    if not completed:
+                        exit_code = 2
+                        print(
+                            "Step",
+                            step.get("code"),
+                            "reported that it hit a timeout and could not complete its work",
+                        )
+
         logger.debug("Finished.")
         sys.exit(exit_code)
     except Exception as ex: