sboms.py

#!/usr/bin/env python3

import json
import os
import subprocess
from datetime import datetime
from contextlib import contextmanager
from pathlib import Path
from typing import Optional, List, Dict

import click
import yardstick
from yardstick.cli import config


SBOM_IMAGE_PREFIX = "ghcr.io/anchore/vml-sbom"
# SBOM_IMAGE_PREFIX = "localhost:5000/anchore/vml-sbom"
DEFAULT_RESULT_SET = "sboms"
TIMESTAMP_OCI_ANNOTATION_KEY = "io.anchore.yardstick.timestamp"
SOURCE_REPO = "https://github.com/anchore/vulnerability-match-labels"
DATA_LICENSE = "CC0-1.0"


@click.option("--verbose", "-v", default=False, help="show logs", is_flag=True)
@click.group(help="Manage SBOMs for the labeled data")
@click.pass_context
def cli(ctx, verbose: bool):
    # pylint: disable=redefined-outer-name, import-outside-toplevel
    import logging.config

    # set the config object to click context to pass to subcommands
    # initialize yardstick based on the current configuration
    ctx.obj: config.Application = config.load()
    yardstick.store.config.set_values(store_root=ctx.obj.store_root)

    log_level = "WARN"
    if verbose:
        log_level = "DEBUG"

    logging.config.dictConfig(
        {
            "version": 1,
            "formatters": {
                "standard": {
                    # [%(module)s.%(funcName)s]
                    "format": "%(asctime)s [%(levelname)s] %(message)s",
                    "datefmt": "",
                },
            },
            "handlers": {
                "default": {
                    "level": log_level,
                    "formatter": "standard",
                    "class": "logging.StreamHandler",
                    "stream": "ext://sys.stderr",
                },
            },
            "loggers": {
                "": {  # root logger
                    "handlers": ["default"],
                    "level": log_level,
                },
            },
        }
    )


@cli.command(name="download", help="download all SBOM blobs from the OCI registry")
@click.option("--result-set", "-r", default=DEFAULT_RESULT_SET, help="the result set to derive which images to download")
@click.pass_obj
def download(cfg: config.Application, result_set: str):
    # 1. derive a set of all images to operate on from the "sbom" result set
    result_set_config = cfg.result_sets.get(result_set, None)
    if not result_set_config:
        raise RuntimeError(f"no result set found for {result_set}")
    
    # 2. download all SBOM blobs from the OCI registry
    scan_requests = result_set_config.scan_requests()
    scan_config_by_oci_ref = {}
    idx = 0
    sbom_scan_requests = [r for r in scan_requests if r.tool.lower().startswith("syft")]
    for idx, r in enumerate(sbom_scan_requests):
        print(f"Pulling SBOM", idx+1, "of", len(sbom_scan_requests), ":", r.image)
        oci_ref, scan_config = download_sbom_results(r)
        if oci_ref and scan_config:
            scan_config_by_oci_ref[oci_ref] = scan_config
        else:
            print("... failed to pull SBOM (skipping)")
        print()

    # 3. update the existing result set with the new downloaded references
    existing_result_set = None
    if yardstick.store.result_set.exists(name=result_set):
        existing_result_set = yardstick.store.result_set.load(result_set)
    else:
        existing_result_set = yardstick.artifact.ResultSet(name=result_set)
        for r in cfg.result_sets[result_set].scan_requests():
            existing_result_set.add(request=r, scan_config=None)

    for state in existing_result_set.state:
        oci_ref = oci_sbom_reference_from_image(state.request.image)
        if oci_ref in scan_config_by_oci_ref:
            state.config = scan_config_by_oci_ref[oci_ref]

    yardstick.store.result_set.save(results=existing_result_set)


def download_sbom_results(request: yardstick.artifact.ScanRequest):
    oci_ref = oci_sbom_reference_from_image(request.image)
    proc = Oras.manifest_fetch(target=oci_ref, capture_output=True, fail_on_error=False)
    if proc.returncode != 0:
        return None, None
    
    manifest = json.loads(proc.stdout)
    annotations = manifest.get("annotations", {})
    timestamp_rfc3339 = annotations.get(TIMESTAMP_OCI_ANNOTATION_KEY, None)
    timestamp = datetime.fromisoformat(timestamp_rfc3339)

    if not timestamp:
        raise RuntimeError(f"no timestamp found for {oci_ref}")
    
    synthesized_config = yardstick.artifact.ScanConfiguration.new(image=request.image, tool=request.tool, timestamp=timestamp)
    data, _ = yardstick.store.scan_result.store_paths(config=synthesized_config)
    parent_dir = os.path.dirname(data)

    if parent_dir and not os.path.exists(parent_dir):
        os.makedirs(parent_dir)

    try:
        Oras.pull(target=oci_ref, destination=parent_dir)
    except RuntimeError:
        return None, None

    scan_config = yardstick.store.scan_result.find_one(by_description=synthesized_config.path)
    if not scan_config:
        raise RuntimeError(f"no scan config found for {synthesized_config.path}")

    return oci_ref, scan_config


@cli.command(name="update", help="upload all SBOM blobs to the OCI registry")
@click.option("--result-set", "-r", default=DEFAULT_RESULT_SET, help="the result set to derive which images to scan")
@click.pass_obj
def update(cfg: config.Application, result_set: str):
    # 1. derive a set of all images to operate on from the "sbom" result set
    result_set_config = cfg.result_sets.get(result_set, None)
    if not result_set_config:
        raise RuntimeError(f"no result set found for {result_set}")

    # 2. run capture on the result set to generate missing SBOMs
    saved_result_set = yardstick.capture.result_set(result_set, result_set_config.scan_requests(), profiles=cfg.profiles.data)


@cli.command(name="upload", help="upload all SBOM blobs to the OCI registry")
@click.option("--result-set", "-r", default=DEFAULT_RESULT_SET, help="the result set to derive which images to scan")
@click.pass_obj
def upload(cfg: config.Application, result_set: str):
    # 1. load existing result set
    saved_result_set = yardstick.store.result_set.load(name=result_set)

    # 2. from the returned result set object, iterate over all SBOMs and upload them to the registry
    for idx, state in enumerate(saved_result_set.state):
        data, _ = yardstick.store.scan_result.store_paths(state.config)
        parent_dir = os.path.dirname(data)

        # note: the SBOM is already in the right format, so we can just upload it
        annotations = {
            TIMESTAMP_OCI_ANNOTATION_KEY: state.config.timestamp_rfc3339,
            "org.opencontainers.image.source": SOURCE_REPO,
            "org.opencontainers.image.licenses": DATA_LICENSE,
            "org.opencontainers.image.description": f"sbom for {state.config.image} captured with yardstick"
        }
        target = oci_sbom_reference_from_image(state.config.image)

        print("Pushing SBOM", idx+1, "of", len(saved_result_set.state), ":", state.request.image)
        Oras.push(target=target, files=["metadata.json", "data.json"], annotations=annotations, cd=parent_dir)
        print()


def oci_sbom_reference_from_image(image: str) -> str:
    img = yardstick.artifact.Image(image)
    # note: we need to keep the docker.io references for stable lookups
    return f"{SBOM_IMAGE_PREFIX}/{img.repository}:sha256-{img.digest.removeprefix('sha256:')}"

def image_from_sbom_reference(oci_sbom_reference: str) -> str:
    return oci_sbom_reference.removeprefix(SBOM_IMAGE_PREFIX+"/").replace(":sha256-", "@sha256:")


@contextmanager
def set_directory(path: Path):
    origin = Path().absolute()
    try:
        os.chdir(path)
        yield
    finally:
        os.chdir(origin)


# why not use the python oras client? 
# The python oras client is not efficient with detecting existing large files (unlike the oras CLI)
class Oras:

    @classmethod
    def pull(cls, target: str, destination: str, **kwargs) -> subprocess.CompletedProcess:
        return cls.run("pull", target, **kwargs, cd=destination)

    @classmethod
    def push(cls, target: str, files: List[str], annotations: Dict[str, str], **kwargs) -> subprocess.CompletedProcess:
        an = []
        for k, v in annotations.items():
            an.append("--annotation")
            an.append(f"{k}={v}")
        return cls.run("push", target, *files, *an, **kwargs)

    @classmethod
    def manifest_fetch(cls, target: str, **kwargs) -> subprocess.CompletedProcess:
        return cls.run("manifest", "fetch", target, **kwargs)


    @classmethod
    def run(cls, *args, cd: Optional[str] = None, fail_on_error: bool = True, **kwargs) -> subprocess.CompletedProcess:
        def call():
            proc = subprocess.run([f"oras", *args], env=os.environ.copy(), **kwargs)

            if fail_on_error and proc.returncode != 0:
                raise RuntimeError(f"return code: {proc.returncode}")

            return proc

        if cd:
            with set_directory(cd):
                return call()
        return call()


if __name__ == '__main__':
    cli()