Azure Omni-Tool MCP Server

import { z } from 'zod'; import { executeAzCommand, executeCommand } from '../lib/cli-executor.js'; import { withCache, CacheKeys } from '../lib/cache.js'; import { createAuditContext } from '../lib/audit.js'; import { getOperatorInfo } from '../lib/config.js'; import { logger } from '../lib/logger.js'; export type QueryType = 'subscriptions' | 'resource_groups' | 'resources' | 'custom'; export const GetAzureContextSchema = z.object({ query_type: z.enum(['subscriptions', 'resource_groups', 'resources', 'custom']).describe('Type of context query'), subscription_id: z.string().optional().describe('Subscription to scope query'), resource_group: z.string().optional().describe('Resource group filter'), custom_query: z.string().optional().describe('KQL query for Resource Graph'), bypass_cache: z.boolean().optional().default(false).describe('Force fresh data'), }); export type GetAzureContextInput = z.infer<typeof GetAzureContextSchema>; export interface ContextResponse { query_type: QueryType; success: boolean; data?: unknown; count?: number; error?: string; command_executed?: string; cached?: boolean; correlation_id?: string; operator?: { email?: string; name?: string }; } type QueryHandler = (input: GetAzureContextInput) => Promise<ContextResponse>; const queryHandlers = new Map<QueryType, QueryHandler>([ ['subscriptions', async ({ bypass_cache }) => { const cmd = 'az account list'; const key = CacheKeys.subscriptions(); const fetcher = async (): Promise<ContextResponse> => { const result = await executeAzCommand(cmd); if (result.success && result.parsedOutput) { const data = result.parsedOutput as unknown[]; return { query_type: 'subscriptions', success: true, data, count: data.length, command_executed: cmd, cached: false }; } return { query_type: 'subscriptions', success: false, error: result.stderr || 'Failed', command_executed: cmd }; }; return bypass_cache ? fetcher() : withCache(key, fetcher); }], ['resource_groups', async ({ subscription_id, bypass_cache }) => { let cmd = 'az group list'; if (subscription_id) cmd += ` --subscription "${subscription_id}"`; const key = CacheKeys.resourceGroups(subscription_id); const fetcher = async (): Promise<ContextResponse> => { const result = await executeAzCommand(cmd); if (result.success && result.parsedOutput) { const data = result.parsedOutput as unknown[]; return { query_type: 'resource_groups', success: true, data, count: data.length, command_executed: cmd, cached: false }; } return { query_type: 'resource_groups', success: false, error: result.stderr || 'Failed', command_executed: cmd }; }; return bypass_cache ? fetcher() : withCache(key, fetcher); }], ['resources', async ({ subscription_id, resource_group, bypass_cache }) => { let cmd = 'az resource list'; if (subscription_id) cmd += ` --subscription "${subscription_id}"`; if (resource_group) cmd += ` --resource-group "${resource_group}"`; const key = CacheKeys.resources(subscription_id, resource_group); const fetcher = async (): Promise<ContextResponse> => { const result = await executeAzCommand(cmd); if (result.success && result.parsedOutput) { const data = result.parsedOutput as unknown[]; return { query_type: 'resources', success: true, data, count: data.length, command_executed: cmd, cached: false }; } return { query_type: 'resources', success: false, error: result.stderr || 'Failed', command_executed: cmd }; }; return bypass_cache ? fetcher() : withCache(key, fetcher); }], ['custom', async ({ custom_query, subscription_id }) => { if (!custom_query) { return { query_type: 'custom', success: false, error: 'custom_query required for type "custom"' }; } const escaped = custom_query.replace(/"/g, '\\"'); let cmd = `az graph query -q "${escaped}"`; if (subscription_id) cmd += ` --subscriptions "${subscription_id}"`; const result = await executeCommand(cmd); if (result.stderr.includes('graph') && result.stderr.includes('not found')) { return { query_type: 'custom', success: false, error: 'Install az extension: az extension add --name resource-graph', command_executed: cmd }; } if (result.success && result.parsedOutput) { const qr = result.parsedOutput as { data?: unknown[]; count?: number }; return { query_type: 'custom', success: true, data: qr.data || qr, count: qr.count ?? (Array.isArray(qr.data) ? qr.data.length : undefined), command_executed: cmd }; } return { query_type: 'custom', success: false, error: result.stderr || 'Query failed', command_executed: cmd }; }], ]); export async function handleGetAzureContext(input: GetAzureContextInput): Promise<ContextResponse> { const { query_type } = input; const operator = getOperatorInfo(); const audit = createAuditContext(`context:${query_type}`, 'low', 'query'); logger.debug('Context query', { query_type, subscription: input.subscription_id, rg: input.resource_group }); const handler = queryHandlers.get(query_type); if (!handler) { return { query_type, success: false, error: `Unknown query type: ${query_type}` }; } const response = await handler(input); response.correlation_id = audit.correlationId; response.operator = operator; if (response.success) { await audit.logSuccess(); logger.info('Context query succeeded', { query_type, count: response.count, cached: response.cached }); } else { await audit.logFailure(response.error || 'Unknown'); logger.warn('Context query failed', { query_type, error: response.error }); } return response; } export const getAzureContextTool = { name: 'get_azure_context', description: `Retrieves Azure environment context. TYPES: subscriptions, resource_groups, resources, custom (KQL) CACHING: 5min default, bypass_cache=true for fresh data EXAMPLES: - Find VMs: custom_query = "Resources | where type == 'microsoft.compute/virtualmachines'" - Find by tag: custom_query = "Resources | where tags.env == 'prod'"`, inputSchema: { type: 'object', properties: { query_type: { type: 'string', enum: ['subscriptions', 'resource_groups', 'resources', 'custom'] }, subscription_id: { type: 'string', description: 'Optional subscription scope' }, resource_group: { type: 'string', description: 'Optional resource group filter' }, custom_query: { type: 'string', description: 'KQL query for Resource Graph' }, bypass_cache: { type: 'boolean', description: 'Force fresh data', default: false }, }, required: ['query_type'], }, };