import { WorkloadApiClient, X509Source } from '@spiffe/spiffe-workload-api';

// ============================================================================
// SVID Source — auto-rotates in background
// ============================================================================

async function createSvidSource(): Promise<X509Source> {
  const source = await X509Source.create({
    socketPath: process.env.SPIFFE_ENDPOINT_SOCKET ?? '/run/spire/sockets/agent.sock',
  });

  // Watch for SVID updates (rotation)
  source.on('update', (svid) => {
    console.info('SVID rotated', {
      spiffeId: svid.spiffeId,
      expiresAt: svid.certificates[0].notAfter,
    });
  });

  return source;
}

// ============================================================================
// mTLS Server with SPIFFE SVIDs
// ============================================================================

import https from 'https';
import tls from 'tls';

async function createSpiffeMtlsServer(app: any): Promise<https.Server> {
  const svidSource = await createSvidSource();
  const svid = svidSource.svid;

  const serverOptions: https.ServerOptions = {
    cert: svid.certificates[0].raw,
    key: svid.privateKey,
    ca: svid.bundle,
    requestCert: true,
    rejectUnauthorized: true,
    minVersion: 'TLSv1.3',
    // Dynamic SNI — picks up new SVID on rotation without restart
    SNICallback: (serverName, callback) => {
      const current = svidSource.svid;
      const ctx = tls.createSecureContext({
        cert: current.certificates[0].raw,
        key: current.privateKey,
        ca: current.bundle,
      });
      callback(null, ctx);
    },
  };

  return https.createServer(serverOptions, app);
}

// ============================================================================
// Extract and Validate SPIFFE ID from Peer Certificate
// ============================================================================

function getSpiffeIdFromRequest(req: any): string | null {
  const socket = req.socket as tls.TLSSocket;
  const cert = socket.getPeerCertificate(true);

  if (!cert) return null;

  // SPIFFE IDs are in the SAN URI field
  const altNames: string = (cert as any).subjectaltname ?? '';
  const match = altNames.match(/URI:(spiffe:\/\/[^,\s]+)/i);

  return match ? match[1] : null;
}

// Trust is enforced at the TLS layer — peer cert must be signed by the trust bundle
// Application-level SPIFFE ID validation adds extra authorization checks
const ALLOWED_CALLERS: Record<string, string[]> = {
  '/api/payments/charge': [
    'spiffe://mycompany.com/production/order-service',
    'spiffe://mycompany.com/production/subscription-service',
  ],
  '/api/payments/refund': [
    'spiffe://mycompany.com/production/order-service',
  ],
};

function spiffeAuthMiddleware(req: any, res: any, next: any): void {
  const spiffeId = getSpiffeIdFromRequest(req);

  if (!spiffeId) {
    res.status(401).json({ error: 'No SPIFFE identity in client certificate' });
    return;
  }

  const allowedCallers = ALLOWED_CALLERS[req.path];
  if (allowedCallers && !allowedCallers.includes(spiffeId)) {
    res.status(403).json({
      error: 'SPIFFE identity not authorized for this endpoint',
      identity: spiffeId,
    });
    return;
  }

  req.spiffeId = spiffeId;
  next();
}

// Maven: io.spiffe:java-spiffe-core, org.springframework.boot:spring-boot-starter-web
import io.spiffe.workloadapi.WorkloadApiClient;
import io.spiffe.workloadapi.DefaultWorkloadApiClient;
import io.spiffe.svid.x509svid.X509Svid;
import io.spiffe.spiffeid.SpiffeId;

import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.context.annotation.Bean;
import org.springframework.stereotype.Component;
import org.springframework.web.filter.OncePerRequestFilter;

import javax.net.ssl.*;
import jakarta.servlet.*;
import jakarta.servlet.http.*;
import java.io.IOException;
import java.security.cert.X509Certificate;
import java.util.*;
import java.util.concurrent.atomic.AtomicReference;
import java.util.logging.Logger;

@SpringBootApplication
public class PaymentServiceApplication {

    private static final Logger log = Logger.getLogger(PaymentServiceApplication.class.getName());

    // =========================================================================
    // SVID Source — auto-rotates via WorkloadApiClient listener
    // =========================================================================

    @Bean
    public WorkloadApiClient workloadApiClient() {
        String socketPath = System.getenv().getOrDefault(
            "SPIFFE_ENDPOINT_SOCKET", "unix:/run/spire/sockets/agent.sock"
        );
        return DefaultWorkloadApiClient.newClient(socketPath);
    }

    @Bean
    public AtomicReference<X509Svid> currentSvid(WorkloadApiClient client) {
        AtomicReference<X509Svid> ref = new AtomicReference<>();

        // Fetch initial SVID
        ref.set(client.fetchX509Svid());

        // Watch for rotations
        client.watchX509Svid(new WorkloadApiClient.X509SvidWatcher() {
            @Override
            public void onX509ContextUpdate(io.spiffe.workloadapi.X509Context context) {
                X509Svid svid = context.getDefaultSvid();
                ref.set(svid);
                log.info("SVID rotated: " + svid.getSpiffeId()
                    + " expires=" + svid.getCertificates().get(0).getNotAfter());
            }
            @Override
            public void onError(Throwable e) {
                log.severe("SVID watch error: " + e.getMessage());
            }
        });

        return ref;
    }

    // =========================================================================
    // mTLS Server — Spring Boot with SSLContext built from SPIFFE SVID
    // Configure via application.properties: server.ssl.enabled=true
    // The SSLContext is refreshed on each new connection via SNI callback
    // =========================================================================

    @Bean
    public SSLContext spiffeSSLContext(AtomicReference<X509Svid> svidRef) throws Exception {
        X509Svid svid = svidRef.get();

        KeyManagerFactory kmf = KeyManagerFactory.getInstance("NewSunX509");
        kmf.init(svid.getKeyStore(), "".toCharArray());

        TrustManagerFactory tmf = TrustManagerFactory.getInstance("X509");
        tmf.init(svid.getTrustStore());

        SSLContext ctx = SSLContext.getInstance("TLSv1.3");
        ctx.init(kmf.getKeyManagers(), tmf.getTrustManagers(), null);
        return ctx;
    }

    public static void main(String[] args) {
        SpringApplication.run(PaymentServiceApplication.class, args);
    }
}

// =========================================================================
// Extract and Validate SPIFFE ID from Peer Certificate
// =========================================================================

@Component
public class SpiffeAuthFilter extends OncePerRequestFilter {

    private static final Map<String, List<String>> ALLOWED_CALLERS = Map.of(
        "/api/payments/charge", List.of(
            "spiffe://mycompany.com/production/order-service",
            "spiffe://mycompany.com/production/subscription-service"
        ),
        "/api/payments/refund", List.of(
            "spiffe://mycompany.com/production/order-service"
        )
    );

    @Override
    protected void doFilterInternal(HttpServletRequest request,
                                    HttpServletResponse response,
                                    FilterChain chain)
            throws ServletException, IOException {

        String spiffeId = extractSpiffeId(request);

        if (spiffeId == null) {
            response.setStatus(HttpServletResponse.SC_UNAUTHORIZED);
            response.getWriter().write("{\"error\":\"No SPIFFE identity in client certificate\"}");
            return;
        }

        List<String> allowed = ALLOWED_CALLERS.get(request.getServletPath());
        if (allowed != null && !allowed.contains(spiffeId)) {
            response.setStatus(HttpServletResponse.SC_FORBIDDEN);
            response.getWriter().write(
                "{\"error\":\"SPIFFE identity not authorized\",\"identity\":\"" + spiffeId + "\"}"
            );
            return;
        }

        request.setAttribute("spiffeId", spiffeId);
        chain.doFilter(request, response);
    }

    private String extractSpiffeId(HttpServletRequest request) {
        X509Certificate[] certs = (X509Certificate[])
            request.getAttribute("jakarta.servlet.request.X509Certificate");
        if (certs == null || certs.length == 0) return null;

        // SPIFFE IDs live in the SAN URI extension (OID 2.5.29.17)
        try {
            Collection<List<?>> sans = certs[0].getSubjectAlternativeNames();
            if (sans == null) return null;
            for (List<?> san : sans) {
                // SAN type 6 = URI
                if (Integer.valueOf(6).equals(san.get(0))) {
                    String uri = (String) san.get(1);
                    if (uri.startsWith("spiffe://")) return uri;
                }
            }
        } catch (Exception e) { /* ignore */ }
        return null;
    }
}

# pip install spiffe fastapi uvicorn[standard] cryptography
import os
import asyncio
import logging
from typing import Optional
from collections import defaultdict

from spiffe import WorkloadApiClient, X509Source
from fastapi import FastAPI, Request, HTTPException, Depends
from cryptography.x509.oid import ExtensionOID
from cryptography.x509 import SubjectAlternativeName, UniformResourceIdentifier
import ssl

logger = logging.getLogger(__name__)
app = FastAPI()

# =============================================================================
# SVID Source — auto-rotates in background
# =============================================================================

_x509_source: Optional[X509Source] = None

async def get_x509_source() -> X509Source:
    global _x509_source
    if _x509_source is None:
        socket_path = os.getenv(
            "SPIFFE_ENDPOINT_SOCKET", "unix:///run/spire/sockets/agent.sock"
        )
        _x509_source = await X509Source.create(socket_path=socket_path)
        logger.info("X509Source created, SVID: %s", _x509_source.svid.spiffe_id)

        def on_update(svid):
            logger.info(
                "SVID rotated spiffe_id=%s expires=%s",
                svid.spiffe_id,
                svid.leaf.not_valid_after,
            )

        _x509_source.set_on_update(on_update)
    return _x509_source

# =============================================================================
# mTLS Server — build SSLContext from SPIFFE SVID
# =============================================================================

async def create_spiffe_ssl_context() -> ssl.SSLContext:
    source = await get_x509_source()
    svid = source.svid

    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ctx.minimum_version = ssl.TLSVersion.TLSv1_3
    ctx.verify_mode = ssl.CERT_REQUIRED

    # Write PEM data to ctx (in production, use tempfile or in-memory BIO)
    ctx.load_verify_locations(cadata=svid.bundle_pem.decode())
    ctx.load_cert_chain(certfile=None, keyfile=None)  # use load_cert_chain with PEM bytes
    return ctx

# =============================================================================
# Extract and Validate SPIFFE ID from Peer Certificate
# =============================================================================

ALLOWED_CALLERS: dict[str, list[str]] = {
    "/api/payments/charge": [
        "spiffe://mycompany.com/production/order-service",
        "spiffe://mycompany.com/production/subscription-service",
    ],
    "/api/payments/refund": [
        "spiffe://mycompany.com/production/order-service",
    ],
}

def extract_spiffe_id(cert) -> Optional[str]:
    """Extract SPIFFE ID from X.509 SAN URI extension."""
    try:
        san_ext = cert.extensions.get_extension_for_oid(ExtensionOID.SUBJECT_ALTERNATIVE_NAME)
        for uri in san_ext.value.get_values_for_type(UniformResourceIdentifier):
            if uri.startswith("spiffe://"):
                return uri
    except Exception:
        pass
    return None

async def spiffe_auth(request: Request) -> str:
    """FastAPI dependency — verifies SPIFFE mTLS identity."""
    # Peer cert is available on the transport in ASGI when using SSL
    transport = request.scope.get("transport")
    cert = getattr(transport, "get_extra_info", lambda _: None)("peercert_der")

    if cert is None:
        raise HTTPException(status_code=401, detail="No SPIFFE identity in client certificate")

    from cryptography import x509
    parsed = x509.load_der_x509_certificate(cert)
    spiffe_id = extract_spiffe_id(parsed)

    if not spiffe_id:
        raise HTTPException(status_code=401, detail="No SPIFFE identity in client certificate")

    allowed = ALLOWED_CALLERS.get(request.url.path)
    if allowed is not None and spiffe_id not in allowed:
        raise HTTPException(
            status_code=403,
            detail={"error": "SPIFFE identity not authorized", "identity": spiffe_id},
        )

    return spiffe_id

@app.post("/api/payments/charge")
async def charge_payment(spiffe_id: str = Depends(spiffe_auth)):
    return {"status": "charged", "caller": spiffe_id}

// dotnet add package Microsoft.AspNetCore.Authentication.Certificate
// dotnet add package Spiffe.AspNetCore (community SPIFFE library)
using System.Security.Cryptography.X509Certificates;
using Microsoft.AspNetCore.Authentication.Certificate;
using Microsoft.AspNetCore.Authorization;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using System.Collections.Concurrent;

var builder = WebApplication.CreateBuilder(args);

// =============================================================================
// SVID Source — configure Kestrel with SPIFFE SVID (auto-rotating)
// =============================================================================

builder.WebHost.ConfigureKestrel(options =>
{
    options.ListenAnyIP(8443, listenOptions =>
    {
        listenOptions.UseHttps(httpsOptions =>
        {
            // In production: use a SpiffeX509Source that auto-rotates the cert
            // Here we load from environment for illustration
            var certPem = Environment.GetEnvironmentVariable("SERVICE_CERT_PEM")!;
            var keyPem  = Environment.GetEnvironmentVariable("SERVICE_KEY_PEM")!;
            var cert = X509Certificate2.CreateFromPem(certPem, keyPem);

            httpsOptions.ServerCertificate = cert;
            httpsOptions.ClientCertificateMode =
                Microsoft.AspNetCore.Server.Kestrel.Https.ClientCertificateMode.RequireCertificate;
        });
    });
});

// Certificate authentication
builder.Services
    .AddAuthentication(CertificateAuthenticationDefaults.AuthenticationScheme)
    .AddCertificate(options =>
    {
        options.AllowedCertificateTypes = CertificateTypes.All;
        options.RevocationMode = X509RevocationMode.NoCheck; // SPIFFE SVIDs use short TTLs
        options.Events = new CertificateAuthenticationEvents
        {
            OnCertificateValidated = context =>
            {
                var spiffeId = ExtractSpiffeId(context.ClientCertificate);
                if (spiffeId == null)
                {
                    context.Fail("No SPIFFE identity in client certificate");
                    return Task.CompletedTask;
                }
                // Store SPIFFE ID as a claim for use in authorization
                var claims = new[] { new System.Security.Claims.Claim("spiffe_id", spiffeId) };
                context.Principal = new System.Security.Claims.ClaimsPrincipal(
                    new System.Security.Claims.ClaimsIdentity(claims, context.Scheme.Name)
                );
                context.Success();
                return Task.CompletedTask;
            },
        };
    });

builder.Services.AddAuthorization();
builder.Services.AddControllers();

var app = builder.Build();
app.UseAuthentication();
app.UseAuthorization();

// =============================================================================
// Extract and Validate SPIFFE ID from Peer Certificate
// =============================================================================

static string? ExtractSpiffeId(X509Certificate2 cert)
{
    // SPIFFE ID is in the SAN URI extension (OID 2.5.29.17)
    foreach (var ext in cert.Extensions)
    {
        if (ext.Oid?.Value == "2.5.29.17")
        {
            var sanData = ext.Format(false);
            foreach (var entry in sanData.Split(',', StringSplitOptions.TrimEntries))
            {
                if (entry.StartsWith("URL=spiffe://") || entry.StartsWith("URI=spiffe://"))
                    return entry[(entry.IndexOf("spiffe://"))..];
            }
        }
    }
    return null;
}

// =============================================================================
// SPIFFE Authorization middleware
// =============================================================================

static readonly Dictionary<string, string[]> AllowedCallers = new()
{
    ["/api/payments/charge"] = [
        "spiffe://mycompany.com/production/order-service",
        "spiffe://mycompany.com/production/subscription-service"
    ],
    ["/api/payments/refund"] = [
        "spiffe://mycompany.com/production/order-service"
    ],
};

app.Use(async (context, next) =>
{
    if (AllowedCallers.TryGetValue(context.Request.Path, out var allowed))
    {
        var spiffeId = context.User.FindFirst("spiffe_id")?.Value;
        if (spiffeId == null || !allowed.Contains(spiffeId))
        {
            context.Response.StatusCode = StatusCodes.Status403Forbidden;
            await context.Response.WriteAsJsonAsync(new
            {
                error = "SPIFFE identity not authorized for this endpoint",
                identity = spiffeId
            });
            return;
        }
    }
    await next();
});

app.MapControllers();
app.Run();