AutopilotMonitor/src/Backend/AutopilotMonitor.Functions/Functions/Vulnerability/SaveCustomCpeMappingFunction.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using System.Text.Json;
using System.Threading.Tasks;
using AutopilotMonitor.Functions.Services;
using AutopilotMonitor.Functions.Services.Vulnerability;
using AutopilotMonitor.Shared;
using Azure.Data.Tables;
using Microsoft.Azure.Functions.Worker;
using Microsoft.Azure.Functions.Worker.Http;
using Microsoft.Extensions.Logging;

namespace AutopilotMonitor.Functions.Functions.Vulnerability
{
    public class SaveCustomCpeMappingFunction
    {
        private readonly TableStorageService _storageService;
        private readonly ILogger<SaveCustomCpeMappingFunction> _logger;

        public SaveCustomCpeMappingFunction(
            TableStorageService storageService,
            ILogger<SaveCustomCpeMappingFunction> logger)
        {
            _storageService = storageService;
            _logger = logger;
        }

        [Function("SaveCustomCpeMapping")]
        public async Task<HttpResponseData> Run(
            [HttpTrigger(AuthorizationLevel.Anonymous, "post", Route = "vulnerability/cpe-mapping")] HttpRequestData req)
        {
            try
            {
                // Galactic Admin only - require Authorization header
                var authHeader = req.Headers.Contains("Authorization")
                    ? req.Headers.GetValues("Authorization").FirstOrDefault()
                    : null;
                if (string.IsNullOrEmpty(authHeader))
                {
                    var unauthorized = req.CreateResponse(HttpStatusCode.Unauthorized);
                    return unauthorized;
                }

                var body = await req.ReadAsStringAsync();
                var mapping = JsonSerializer.Deserialize<CpeMappingRequest>(body ?? "",
                    new JsonSerializerOptions { PropertyNameCaseInsensitive = true });

                if (mapping == null || string.IsNullOrWhiteSpace(mapping.NormalizedProduct) || string.IsNullOrWhiteSpace(mapping.CpeUri))
                {
                    var badReq = req.CreateResponse(HttpStatusCode.BadRequest);
                    await badReq.WriteAsJsonAsync(new { success = false, message = "normalizedProduct and cpeUri are required" });
                    return badReq;
                }

                var tableClient = _storageService.GetTableClient(Constants.TableNames.VulnerabilityCache);

                var rowKey = $"{(mapping.NormalizedVendor ?? "unknown")}:{mapping.NormalizedProduct}".ToLowerInvariant();
                // Sanitize RowKey (Table Storage doesn't allow /, \, #, ?)
                rowKey = rowKey.Replace("/", "_").Replace("\\", "_").Replace("#", "_").Replace("?", "_");

                var entity = new TableEntity("cpe_map_custom", rowKey)
                {
                    { "NormalizedVendor", mapping.NormalizedVendor ?? "" },
                    { "NormalizedProduct", mapping.NormalizedProduct },
                    { "CpeVendor", mapping.CpeVendor ?? "" },
                    { "CpeProduct", mapping.CpeProduct ?? "" },
                    { "CpeUri", mapping.CpeUri },
                    { "Category", mapping.Category ?? "custom" },
                    { "DisplayNamePatternsJson", JsonSerializer.Serialize(mapping.DisplayNamePatterns ?? new List<string>()) },
                    { "PublisherPatternsJson", JsonSerializer.Serialize(mapping.PublisherPatterns ?? new List<string>()) },
                    { "Source", "manual" },
                    { "CreatedAt", DateTime.UtcNow.ToString("o") }
                };

                await tableClient.UpsertEntityAsync(entity);

                // Reset the static mappings cache so future correlations use the new mapping
                VulnerabilityCorrelationService.ResetMappingsCache();

                // Update SoftwareInventory entries that match this product so they no longer appear as "unmatched"
                var updatedInventoryCount = await UpdateSoftwareInventoryCpeAsync(
                    mapping.NormalizedProduct, mapping.CpeUri);

                _logger.LogInformation(
                    "Custom CPE mapping saved: {RowKey} -> {CpeUri}, updated {Count} inventory entries",
                    rowKey, mapping.CpeUri, updatedInventoryCount);

                var response = req.CreateResponse(HttpStatusCode.OK);
                await response.WriteAsJsonAsync(new { success = true, message = "CPE mapping saved", rowKey });
                return response;
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Failed to save custom CPE mapping");
                var errorResponse = req.CreateResponse(HttpStatusCode.InternalServerError);
                await errorResponse.WriteAsJsonAsync(new { success = false, message = "Failed to save mapping" });
                return errorResponse;
            }
        }

        /// <summary>
        /// Updates SoftwareInventory entries whose NormalizedName matches the mapped product,
        /// setting their CpeUri so they no longer appear in the unmatched software list.
        /// </summary>
        private async Task<int> UpdateSoftwareInventoryCpeAsync(string normalizedProduct, string cpeUri)
        {
            try
            {
                var inventoryClient = _storageService.GetTableClient(Constants.TableNames.SoftwareInventory);
                var allEntities = inventoryClient.QueryAsync<TableEntity>(maxPerPage: 1000);
                var updatedCount = 0;

                await foreach (var entity in allEntities)
                {
                    var existingCpe = entity.GetString("CpeUri") ?? "";
                    if (!string.IsNullOrEmpty(existingCpe))
                        continue; // Already has a CPE mapping

                    var normalizedName = (entity.GetString("NormalizedName") ?? "").ToLowerInvariant();
                    if (normalizedName == normalizedProduct.ToLowerInvariant())
                    {
                        entity["CpeUri"] = cpeUri;
                        await inventoryClient.UpsertEntityAsync(entity);
                        updatedCount++;
                    }
                }

                return updatedCount;
            }
            catch (Exception ex)
            {
                _logger.LogWarning(ex, "Failed to update SoftwareInventory entries for {Product}", normalizedProduct);
                return 0;
            }
        }

        private class CpeMappingRequest
        {
            public string NormalizedVendor { get; set; } = "";
            public string NormalizedProduct { get; set; } = "";
            public string CpeVendor { get; set; } = "";
            public string CpeProduct { get; set; } = "";
            public string CpeUri { get; set; } = "";
            public string Category { get; set; } = "custom";
            public List<string> DisplayNamePatterns { get; set; } = new();
            public List<string> PublisherPatterns { get; set; } = new();
        }
    }
}