Umbrella MCP Server

import { SignJWT, jwtVerify, exportJWK, JWK } from 'jose'; import { readFileSync, writeFileSync, existsSync, mkdirSync } from 'fs'; import { resolve, join } from 'path'; import crypto from 'crypto'; import { ConfigLoader } from '../config/config-loader.js'; import { AuthToken } from '../types.js'; import type { UserManagementInfo } from '../dual-auth.js'; type KeyLike = any; // jose v6 doesn't export KeyLike directly export interface TokenPayload { sub: string; // User subject (email or ID) clientId: string; // Tenant ID umbrellaAuth: AuthToken; // Direct Umbrella auth token (includes userManagementInfo) exp: number; iat: number; } export class TokenManager { private config = ConfigLoader.getInstance(); private privateKey: KeyLike | null = null; private publicKey: KeyLike | null = null; private jwkPublic: JWK | null = null; constructor() { this.initializeRSAKeys(); } private async initializeRSAKeys(): Promise<void> { const authConfig = this.config.getConfig().auth; const keyPath = resolve(authConfig.rsaKeyPath); // Create keys directory if it doesn't exist if (!existsSync(keyPath)) { mkdirSync(keyPath, { recursive: true }); } const privateKeyPath = join(keyPath, 'private.pem'); const publicKeyPath = join(keyPath, 'public.pem'); const jwkPath = join(keyPath, 'jwk.json'); try { // Try to load existing keys if (existsSync(privateKeyPath) && existsSync(publicKeyPath)) { const privateKeyPem = readFileSync(privateKeyPath, 'utf-8'); const publicKeyPem = readFileSync(publicKeyPath, 'utf-8'); this.privateKey = crypto.createPrivateKey(privateKeyPem); this.publicKey = crypto.createPublicKey(publicKeyPem); if (existsSync(jwkPath)) { this.jwkPublic = JSON.parse(readFileSync(jwkPath, 'utf-8')); } else { this.jwkPublic = await exportJWK(this.publicKey); this.jwkPublic.kid = this.generateKid(this.publicKey); writeFileSync(jwkPath, JSON.stringify(this.jwkPublic, null, 2)); } console.log('[TOKEN] Loaded existing RSA keys'); } else { // Generate new keys await this.generateAndSaveKeys(privateKeyPath, publicKeyPath, jwkPath); } } catch (error) { console.error('[TOKEN] Error loading keys, generating new ones:', error); await this.generateAndSaveKeys(privateKeyPath, publicKeyPath, jwkPath); } } private async generateAndSaveKeys(privateKeyPath: string, publicKeyPath: string, jwkPath: string): Promise<void> { console.log('[TOKEN] Generating new RSA key pair...'); // Use Node's crypto to generate RSA keys const { privateKey, publicKey } = crypto.generateKeyPairSync('rsa', { modulusLength: 2048, publicKeyEncoding: { type: 'spki', format: 'pem' }, privateKeyEncoding: { type: 'pkcs8', format: 'pem' } }); // Store as strings first const privateKeyPem = privateKey; const publicKeyPem = publicKey; // Save to files writeFileSync(privateKeyPath, privateKeyPem); writeFileSync(publicKeyPath, publicKeyPem); // Convert to KeyLike for jose this.privateKey = crypto.createPrivateKey(privateKeyPem); this.publicKey = crypto.createPublicKey(publicKeyPem); // Export and save JWK this.jwkPublic = await exportJWK(this.publicKey); this.jwkPublic.kid = this.generateKid(this.publicKey); writeFileSync(jwkPath, JSON.stringify(this.jwkPublic, null, 2)); console.log('[TOKEN] RSA keys generated and saved'); } private generateKid(publicKey: KeyLike): string { // Generate a deterministic key ID from the public key const publicKeyPem = publicKey.export({ type: 'spki', format: 'pem' }) as string; return crypto.createHash('sha256').update(publicKeyPem).digest('hex').substring(0, 8); } /** * Create a JWT containing Umbrella auth token directly (no encryption) */ async createToken(userEmail: string, clientId: string, umbrellaAuth: AuthToken): Promise<string> { if (!this.privateKey) { await this.initializeRSAKeys(); } const authConfig = this.config.getConfig().auth; const now = Math.floor(Date.now() / 1000); const exp = now + authConfig.tokenExpiration; // Simply embed the Umbrella auth token directly in the JWT (including userManagementInfo) const token = await new SignJWT({ sub: userEmail, clientId, umbrellaAuth: { Authorization: umbrellaAuth.Authorization, userManagementInfo: umbrellaAuth.userManagementInfo } }) .setProtectedHeader({ alg: 'RS256', kid: this.jwkPublic?.kid || 'default' }) .setIssuer(authConfig.issuer) .setIssuedAt(now) .setExpirationTime(exp) .sign(this.privateKey!); if (this.config.isDebugMode()) { console.log('[TOKEN] Created token for:', { userEmail, clientId, expiresIn: `${authConfig.tokenExpiration}s` }); } return token; } /** * Verify token and extract Umbrella credentials directly */ async verifyToken(token: string): Promise<{ payload: TokenPayload; umbrellaAuth: AuthToken }> { if (!this.publicKey) { await this.initializeRSAKeys(); } const authConfig = this.config.getConfig().auth; const { payload } = await jwtVerify(token, this.publicKey!, { issuer: authConfig.issuer, algorithms: ['RS256'] }); const tokenPayload = payload as unknown as TokenPayload; // Extract Umbrella auth directly (no decryption needed), including userManagementInfo const umbrellaAuth: AuthToken = { Authorization: tokenPayload.umbrellaAuth.Authorization, userManagementInfo: tokenPayload.umbrellaAuth.userManagementInfo }; return { payload: tokenPayload, umbrellaAuth }; } /** * Check if token is expired */ isTokenExpired(exp: number): boolean { const now = Math.floor(Date.now() / 1000); return now >= exp; } /** * Get JWK for public key (used by JWKS endpoint) */ getPublicJWK(): JWK | null { return this.jwkPublic; } }