git.stella-ops.org/src/AdvisoryAI/StellaOps.AdvisoryAI/Retrievers/SbomContextRetriever.cs

using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging;
using StellaOps.AdvisoryAI.Abstractions;
using StellaOps.AdvisoryAI.Context;
using StellaOps.AdvisoryAI.Providers;

namespace StellaOps.AdvisoryAI.Retrievers;

internal sealed class SbomContextRetriever : ISbomContextRetriever
{
    private readonly ISbomContextClient _client;
    private readonly ILogger<SbomContextRetriever>? _logger;

    public SbomContextRetriever(ISbomContextClient client, ILogger<SbomContextRetriever>? logger = null)
    {
        _client = client ?? throw new ArgumentNullException(nameof(client));
        _logger = logger;
    }

    public async Task<SbomContextResult> RetrieveAsync(SbomContextRequest request, CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(request);

        var query = new SbomContextQuery(
            request.ArtifactId,
            request.Purl,
            request.MaxTimelineEntries,
            request.MaxDependencyPaths,
            request.IncludeEnvironmentFlags,
            request.IncludeBlastRadius);

        SbomContextDocument? document;
        try
        {
            document = await _client.GetContextAsync(query, cancellationToken).ConfigureAwait(false);
        }
        catch (Exception ex)
        {
            _logger?.LogError(ex, "Failed to retrieve SBOM context for artifact {ArtifactId}", request.ArtifactId);
            document = null;
        }

        if (document is null)
        {
            _logger?.LogWarning("No SBOM context returned for artifact {ArtifactId}", request.ArtifactId);
            return SbomContextResult.Empty(request.ArtifactId, request.Purl);
        }

        var timeline = ShapeTimeline(document.Versions, request.MaxTimelineEntries);
        var paths = ShapeDependencyPaths(document.DependencyPaths, request.MaxDependencyPaths);
        var environmentFlags = request.IncludeEnvironmentFlags
            ? NormalizeEnvironmentFlags(document.EnvironmentFlags)
            : ImmutableDictionary<string, string>.Empty;
        var blastRadius = request.IncludeBlastRadius ? ShapeBlastRadius(document.BlastRadius) : null;
        var metadata = BuildMetadata(document, timeline.Count, paths.Count, environmentFlags.Count, blastRadius is not null);

        return SbomContextResult.Create(
            document.ArtifactId,
            document.Purl,
            timeline,
            paths,
            environmentFlags,
            blastRadius,
            metadata);
    }

    private static IReadOnlyList<SbomVersionTimelineEntry> ShapeTimeline(ImmutableArray<SbomVersionRecord> versions, int max)
    {
        if (versions.IsDefaultOrEmpty || max == 0)
        {
            return Array.Empty<SbomVersionTimelineEntry>();
        }

        return versions
            .OrderBy(static v => v.FirstObserved)
            .ThenBy(static v => v.Version, StringComparer.Ordinal)
            .Take(max > 0 ? max : int.MaxValue)
            .Select(static v => new SbomVersionTimelineEntry(
                v.Version,
                v.FirstObserved,
                v.LastObserved,
                string.IsNullOrWhiteSpace(v.Status)
                    ? (v.IsFixAvailable ? "fixed" : "unknown")
                    : v.Status.Trim(),
                string.IsNullOrWhiteSpace(v.Source) ? "sbom" : v.Source.Trim()))
            .ToImmutableArray();
    }

    private static IReadOnlyList<SbomDependencyPath> ShapeDependencyPaths(ImmutableArray<SbomDependencyPathRecord> paths, int max)
    {
        if (paths.IsDefaultOrEmpty || max == 0)
        {
            return Array.Empty<SbomDependencyPath>();
        }

        var distinct = new SortedDictionary<string, SbomDependencyPath>(StringComparer.Ordinal);
        foreach (var path in paths)
        {
            if (path.Nodes.IsDefaultOrEmpty)
            {
                continue;
            }

            var nodeList = path.Nodes
                .Select(static node => new SbomDependencyNode(node.Identifier, node.Version))
                .ToImmutableArray();

            if (nodeList.IsDefaultOrEmpty)
            {
                continue;
            }

            var key = string.Join(
                "|",
                nodeList.Select(static n => string.Concat(n.Identifier, "@", n.Version ?? string.Empty)));

            if (distinct.ContainsKey(key))
            {
                continue;
            }

            var dependencyPath = new SbomDependencyPath(
                nodeList,
                path.IsRuntime,
                string.IsNullOrWhiteSpace(path.Source) ? null : path.Source.Trim());

            distinct[key] = dependencyPath;
        }

        return distinct.Values
            .OrderBy(p => p.IsRuntime ? 0 : 1)
            .ThenBy(p => p.Nodes.Length)
            .ThenBy(p => string.Join("|", p.Nodes.Select(n => n.Identifier)), StringComparer.Ordinal)
            .Take(max > 0 ? max : int.MaxValue)
            .ToImmutableArray();
    }

    private static IReadOnlyDictionary<string, string> NormalizeEnvironmentFlags(ImmutableDictionary<string, string> flags)
    {
        if (flags == default || flags.IsEmpty)
        {
            return ImmutableDictionary<string, string>.Empty;
        }

        var builder = ImmutableSortedDictionary.CreateBuilder<string, string>(StringComparer.Ordinal);
        foreach (var pair in flags)
        {
            if (string.IsNullOrWhiteSpace(pair.Key) || pair.Value is null)
            {
                continue;
            }

            builder[pair.Key.Trim()] = pair.Value.Trim();
        }

        return builder.ToImmutable();
    }

    private static SbomBlastRadiusSummary? ShapeBlastRadius(SbomBlastRadiusRecord? record)
    {
        if (record is null)
        {
            return null;
        }

        var metadata = record.Metadata == default
            ? ImmutableDictionary<string, string>.Empty
            : record.Metadata.ToImmutableDictionary(StringComparer.Ordinal);

        return new SbomBlastRadiusSummary(
            record.ImpactedAssets,
            record.ImpactedWorkloads,
            record.ImpactedNamespaces,
            record.ImpactedPercentage,
            metadata);
    }

    private static IReadOnlyDictionary<string, string> BuildMetadata(
        SbomContextDocument document,
        int timelineCount,
        int pathCount,
        int environmentFlagCount,
        bool hasBlastRadius)
    {
        var builder = ImmutableDictionary.CreateBuilder<string, string>(StringComparer.Ordinal);
        foreach (var pair in document.Metadata)
        {
            if (string.IsNullOrWhiteSpace(pair.Key) || pair.Value is null)
            {
                continue;
            }

            builder[pair.Key.Trim()] = pair.Value.Trim();
        }

        builder["version_count"] = timelineCount.ToString();
        builder["dependency_path_count"] = pathCount.ToString();
        builder["environment_flag_count"] = environmentFlagCount.ToString();
        builder["blast_radius_present"] = hasBlastRadius.ToString();

        return builder.ToImmutable();
    }
}