sh-mcp/infra/lib/auth-stack.ts

/**
 * sh-mcp-auth — the Cognito auth substrate (design.md §2; agentforce-plan §0.1).
 *
 * Stands up the OAuth2.1 broker the servers already validate against: a Google-
 * federated Cognito user pool (ESSENTIALS feature plan — required for the V2
 * pre-token trigger), per-tier app clients whose `AllowedOAuthScopes` are the
 * PRIMARY trust-tier boundary, a SUPPRESS-ONLY pre-token Lambda, a 5-minute
 * group-sync Lambda, and the sync-state + deny-list tables.
 *
 * Surface-agnostic: app-client callback URLs come from CDK context
 * (`callbackUrls`), defaulting to a placeholder until the Agentforce-vs-Bolt
 * surface is chosen. No API Gateway / WAF / server hosting here (Phase 2b).
 */

import {
  Stack,
  type StackProps,
  Duration,
  RemovalPolicy,
  SecretValue,
  aws_cognito as cognito,
  aws_dynamodb as dynamodb,
  aws_kms as kms,
  aws_lambda as lambda,
  aws_logs as logs,
  aws_iam as iam,
  aws_sns as sns,
  aws_events as events,
  aws_events_targets as targets,
  aws_cloudwatch as cw,
  aws_cloudwatch_actions as cwactions,
} from 'aws-cdk-lib';
import { NodejsFunction } from 'aws-cdk-lib/aws-lambda-nodejs';
import type { Construct } from 'constructs';
import * as path from 'node:path';
import { fileURLToPath } from 'node:url';

const __dirname = path.dirname(fileURLToPath(import.meta.url));
const repoRoot = path.resolve(__dirname, '..', '..');

// Shared account CMKs (referenced by ARN, NOT fromLookup, so synth needs no AWS
// creds — memory: alarm-topic + sensitive-log-group CMK conventions).
const ACCOUNT = '328440206208';
const REGION = 'us-east-1';
const ALARM_KMS_ARN = `arn:aws:kms:${REGION}:${ACCOUNT}:key/abad16b5-5474-43b9-bbc4-89fc8a8a6ecf`; // alias/seahaven-alarm-topics

/** Resource-server-prefixed scope strings, kept in sync with the pre-token Lambda. */
const OPS_SCOPES = ['ops:read', 'ops:tasks', 'gmail:self', 'calendar:self'];
const FINANCE_SCOPES = ['finance:read', 'finance:admin'];

export class ShMcpAuthStack extends Stack {
  constructor(scope: Construct, id: string, props?: StackProps) {
    super(scope, id, props);

    const alarmKey = kms.Key.fromKeyArn(this, 'AlarmKey', ALARM_KMS_ARN);
    // alias/seahaven-logs CMK ARN is supplied at deploy via context (-c logsKmsArn=...)
    // to avoid hardcoding/guessing the key id; these log groups carry no secrets
    // (scope decisions + counts only), so absent context we use AWS-managed encryption.
    const logsKmsArn = this.node.tryGetContext('logsKmsArn') as string | undefined;
    const logsKey = logsKmsArn ? kms.Key.fromKeyArn(this, 'LogsKey', logsKmsArn) : undefined;

    // --- Alarm topic (CMK alias/seahaven-alarm-topics; never alias/aws/sns) ---
    const alarmTopic = new sns.Topic(this, 'AlarmTopic', {
      topicName: 'sh-mcp-alarms',
      masterKey: alarmKey,
    });

    // --- DynamoDB: sync-state freshness marker + deny-list (hard revocation) ---
    const syncState = new dynamodb.Table(this, 'SyncState', {
      tableName: 'sh-mcp-sync-state',
      partitionKey: { name: 'pk', type: dynamodb.AttributeType.STRING },
      billingMode: dynamodb.BillingMode.PAY_PER_REQUEST,
      encryption: dynamodb.TableEncryption.AWS_MANAGED,
      pointInTimeRecoverySpecification: { pointInTimeRecoveryEnabled: true },
      removalPolicy: RemovalPolicy.RETAIN,
    });
    // Stable logical IDs so a future construct-path refactor never replaces (and
    // drops) these tables (handbook: never remove overrideLogicalId).
    (syncState.node.defaultChild as dynamodb.CfnTable).overrideLogicalId('SyncStateTable');

    const denyList = new dynamodb.Table(this, 'DenyList', {
      tableName: 'sh-mcp-deny-list',
      partitionKey: { name: 'sub', type: dynamodb.AttributeType.STRING },
      billingMode: dynamodb.BillingMode.PAY_PER_REQUEST,
      encryption: dynamodb.TableEncryption.AWS_MANAGED,
      pointInTimeRecoverySpecification: { pointInTimeRecoveryEnabled: true },
      timeToLiveAttribute: 'expiresAt', // epoch seconds; auto-expires a revocation
      removalPolicy: RemovalPolicy.RETAIN,
    });
    (denyList.node.defaultChild as dynamodb.CfnTable).overrideLogicalId('DenyListTable');

    // --- Lambdas (Node, arm64, explicit 60-day CMK-encrypted log groups) ---
    const preTokenLogs = new logs.LogGroup(this, 'PreTokenLogs', {
      logGroupName: '/aws/lambda/sh-mcp-pre-token-gen',
      retention: logs.RetentionDays.TWO_MONTHS,
      encryptionKey: logsKey,
      removalPolicy: RemovalPolicy.RETAIN,
    });
    const preTokenFn = new NodejsFunction(this, 'PreTokenFn', {
      functionName: 'sh-mcp-pre-token-gen',
      runtime: lambda.Runtime.NODEJS_22_X,
      architecture: lambda.Architecture.ARM_64,
      entry: path.join(repoRoot, 'auth', 'pre-token-gen', 'src', 'index.ts'),
      handler: 'handler',
      timeout: Duration.seconds(5),
      memorySize: 256,
      logGroup: preTokenLogs,
      environment: {
        SYNC_STATE_TABLE: syncState.tableName,
        DENY_LIST_TABLE: denyList.tableName,
      },
    });
    // Least privilege: read-only on the sync-state marker and the deny-list
    // (consulted at every mint for hard revocation), nothing else.
    syncState.grantReadData(preTokenFn);
    denyList.grantReadData(preTokenFn);

    const groupSyncLogs = new logs.LogGroup(this, 'GroupSyncLogs', {
      logGroupName: '/aws/lambda/sh-mcp-group-sync',
      retention: logs.RetentionDays.TWO_MONTHS,
      encryptionKey: logsKey,
      removalPolicy: RemovalPolicy.RETAIN,
    });
    const groupSyncFn = new NodejsFunction(this, 'GroupSyncFn', {
      functionName: 'sh-mcp-group-sync',
      runtime: lambda.Runtime.NODEJS_22_X,
      architecture: lambda.Architecture.ARM_64,
      entry: path.join(repoRoot, 'auth', 'group-sync', 'src', 'index.ts'),
      handler: 'handler',
      timeout: Duration.seconds(60),
      memorySize: 256,
      logGroup: groupSyncLogs,
      environment: {
        USER_POOL_ID: '', // set below once the pool exists (avoids a cycle)
        SYNC_STATE_TABLE: syncState.tableName,
        GOOGLE_SA_SECRET_ARN: `arn:aws:secretsmanager:${REGION}:${ACCOUNT}:secret:sh-mcp/google-directory-sa`,
        WORKSPACE_DOMAIN: 'seahavenind.com',
      },
    });
    syncState.grantWriteData(groupSyncFn);
    // Scoped Secrets Manager read for ONLY the Google service-account secret.
    groupSyncFn.addToRolePolicy(
      new iam.PolicyStatement({
        actions: ['secretsmanager:GetSecretValue'],
        resources: [
          `arn:aws:secretsmanager:${REGION}:${ACCOUNT}:secret:sh-mcp/google-directory-sa-*`,
        ],
      }),
    );

    // --- Cognito user pool (ESSENTIALS — required for the V2 pre-token trigger) ---
    const userPool = new cognito.UserPool(this, 'UserPool', {
      userPoolName: 'sh-mcp',
      featurePlan: cognito.FeaturePlan.ESSENTIALS,
      selfSignUpEnabled: false,
      signInAliases: { email: true },
      standardAttributes: { email: { required: true, mutable: true } },
      removalPolicy: RemovalPolicy.RETAIN,
    });

    // Google as an external OIDC IdP. client_id/secret resolve from Secrets
    // Manager at deploy via a CFN dynamic reference (never inlined in the template).
    const googleIdp = new cognito.CfnUserPoolIdentityProvider(this, 'GoogleIdp', {
      userPoolId: userPool.userPoolId,
      providerName: 'Google',
      providerType: 'Google',
      attributeMapping: { email: 'email', username: 'sub' },
      providerDetails: {
        client_id: SecretValue.secretsManager('sh-mcp/google-oidc', {
          jsonField: 'client_id',
        }).toString(),
        client_secret: SecretValue.secretsManager('sh-mcp/google-oidc', {
          jsonField: 'client_secret',
        }).toString(),
        authorize_scopes: 'openid email profile',
      },
    });

    // Resource servers carry the custom scopes the access tokens are prefixed with.
    const opsRs = userPool.addResourceServer('OpsResourceServer', {
      identifier: 'sh-mcp-ops',
      scopes: OPS_SCOPES.map(
        (s) => new cognito.ResourceServerScope({ scopeName: s, scopeDescription: s }),
      ),
    });
    const financeRs = userPool.addResourceServer('FinanceResourceServer', {
      identifier: 'sh-mcp-finance',
      scopes: FINANCE_SCOPES.map(
        (s) => new cognito.ResourceServerScope({ scopeName: s, scopeDescription: s }),
      ),
    });

    const rsScope = (rs: cognito.IUserPoolResourceServer, name: string): cognito.OAuthScope =>
      cognito.OAuthScope.resourceServer(
        rs,
        new cognito.ResourceServerScope({ scopeName: name, scopeDescription: name }),
      );

    const callbackUrls = (this.node.tryGetContext('callbackUrls') as string[] | undefined) ?? [
      'https://placeholder.seahavenind.com/oauth/callback', // surface TBD (Agentforce vs Bolt)
    ];

    // Per-tier app clients. AllowedOAuthScopes IS the trust-tier ceiling: finance
    // is finance:read ONLY; exec is ops + gmail/calendar with NO finance; ops is
    // read+tasks. A client physically cannot request a scope outside its tier.
    const mkClient = (
      cid: string,
      scopes: cognito.OAuthScope[],
      accessTokenValidity: Duration,
      // Per-tier refresh window. The refresh token is the real persistence horizon:
      // a 15-min access TTL is meaningless if a stolen refresh token re-mints for
      // 30 days, so the sensitive finance tier gets a short window of its own.
      refreshTokenValidity: Duration,
    ): cognito.UserPoolClient =>
      userPool.addClient(cid, {
        userPoolClientName: cid,
        generateSecret: true,
        oAuth: {
          flows: { authorizationCodeGrant: true },
          scopes: [cognito.OAuthScope.OPENID, cognito.OAuthScope.EMAIL, ...scopes],
          callbackUrls,
        },
        supportedIdentityProviders: [cognito.UserPoolClientIdentityProvider.GOOGLE],
        accessTokenValidity,
        idTokenValidity: accessTokenValidity,
        refreshTokenValidity,
        enableTokenRevocation: true,
        preventUserExistenceErrors: true,
      });

    const opsClient = mkClient(
      'sh-agentforce-ops',
      [rsScope(opsRs, 'ops:read'), rsScope(opsRs, 'ops:tasks')],
      Duration.minutes(60),
      Duration.days(30),
    );
    const execClient = mkClient(
      'sh-agentforce-exec',
      [
        rsScope(opsRs, 'ops:read'),
        rsScope(opsRs, 'ops:tasks'),
        rsScope(opsRs, 'gmail:self'),
        rsScope(opsRs, 'calendar:self'),
      ],
      Duration.minutes(60),
      Duration.days(30),
    );
    const financeClient = mkClient(
      'sh-agentforce-finance',
      [rsScope(financeRs, 'finance:read')], // finance:read ONLY — 15-min access TTL
      Duration.minutes(15),
      Duration.hours(8), // short refresh horizon: a stolen finance refresh token dies in 8h, not 30d
    );
    for (const c of [opsClient, execClient, financeClient]) c.node.addDependency(googleIdp);

    // Cognito groups the group-sync Lambda reconciles from Google Groups.
    for (const g of ['sh-mcp-ops', 'sh-mcp-assistant', 'sh-mcp-finance', 'sh-mcp-admin']) {
      new cognito.CfnUserPoolGroup(this, `Group-${g}`, {
        userPoolId: userPool.userPoolId,
        groupName: g,
      });
    }

    // Wire the group-sync Lambda's pool id now that the pool exists, and grant
    // least-privilege Cognito group-management on THIS pool only.
    groupSyncFn.addEnvironment('USER_POOL_ID', userPool.userPoolId);
    groupSyncFn.addToRolePolicy(
      new iam.PolicyStatement({
        actions: [
          'cognito-idp:CreateGroup',
          'cognito-idp:ListUsersInGroup',
          'cognito-idp:ListUsers',
          'cognito-idp:AdminAddUserToGroup',
          'cognito-idp:AdminRemoveUserFromGroup',
        ],
        resources: [userPool.userPoolArn],
      }),
    );

    // Attach the SUPPRESS-ONLY pre-token Lambda as a V2 trigger (escape hatch:
    // the L2 addTrigger does not expose LambdaVersion, and V2 is required for
    // scope override). Grant Cognito permission to invoke it.
    const cfnPool = userPool.node.defaultChild as cognito.CfnUserPool;
    cfnPool.lambdaConfig = {
      preTokenGenerationConfig: {
        lambdaArn: preTokenFn.functionArn,
        lambdaVersion: 'V2_0',
      },
    };
    preTokenFn.addPermission('CognitoInvoke', {
      principal: new iam.ServicePrincipal('cognito-idp.amazonaws.com'),
      sourceArn: userPool.userPoolArn,
    });

    // --- Group-sync schedule: every 5 minutes ---
    new events.Rule(this, 'GroupSyncSchedule', {
      ruleName: 'sh-mcp-group-sync',
      schedule: events.Schedule.rate(Duration.minutes(5)),
      targets: [new targets.LambdaFunction(groupSyncFn)],
    });

    // --- Alarms (ALARM-state actions only, CMK topic) ---
    const onAlarm = new cwactions.SnsAction(alarmTopic);

    // Group-sync failure (any error in a 15-min window).
    const groupSyncErrors = groupSyncFn
      .metricErrors({ period: Duration.minutes(15), statistic: 'Sum' })
      .createAlarm(this, 'GroupSyncErrorsAlarm', {
        alarmName: 'sh-mcp-group-sync-errors',
        threshold: 1,
        evaluationPeriods: 1,
        comparisonOperator: cw.ComparisonOperator.GREATER_THAN_OR_EQUAL_TO_THRESHOLD,
        treatMissingData: cw.TreatMissingData.NOT_BREACHING,
      });
    groupSyncErrors.addAlarmAction(onAlarm);

    // Group-sync stopped running entirely (two-alarm pattern for the 5-min job).
    const groupSyncMissing = groupSyncFn
      .metricInvocations({ period: Duration.minutes(15), statistic: 'Sum' })
      .createAlarm(this, 'GroupSyncMissingAlarm', {
        alarmName: 'sh-mcp-group-sync-missing',
        threshold: 1,
        evaluationPeriods: 1,
        comparisonOperator: cw.ComparisonOperator.LESS_THAN_THRESHOLD,
        treatMissingData: cw.TreatMissingData.BREACHING,
      });
    groupSyncMissing.addAlarmAction(onAlarm);

    // Pre-token Lambda error rate (a bug here blocks or mis-scopes token issuance).
    const preTokenErrors = preTokenFn
      .metricErrors({ period: Duration.minutes(5), statistic: 'Sum' })
      .createAlarm(this, 'PreTokenErrorsAlarm', {
        alarmName: 'sh-mcp-pre-token-errors',
        threshold: 1,
        evaluationPeriods: 1,
        comparisonOperator: cw.ComparisonOperator.GREATER_THAN_OR_EQUAL_TO_THRESHOLD,
        treatMissingData: cw.TreatMissingData.NOT_BREACHING,
      });
    preTokenErrors.addAlarmAction(onAlarm);

    // Expose ids for the servers' config (consumed in Phase 2b wiring).
    this.exportValue(userPool.userPoolId, { name: 'sh-mcp-user-pool-id' });
    this.exportValue(opsClient.userPoolClientId, { name: 'sh-mcp-ops-client-id' });
    this.exportValue(execClient.userPoolClientId, { name: 'sh-mcp-exec-client-id' });
    this.exportValue(financeClient.userPoolClientId, { name: 'sh-mcp-finance-client-id' });
  }
}