sh-mcp/auth/group-sync/src/aws.ts

/**
 * AWS-backed implementations of the group-sync target + sync-state writer.
 * Externalized at bundle time; only exercised in deployed (aws) runs.
 */

import {
  CognitoIdentityProviderClient,
  CreateGroupCommand,
  ListUsersInGroupCommand,
  ListUsersCommand,
  AdminAddUserToGroupCommand,
  AdminRemoveUserFromGroupCommand,
} from '@aws-sdk/client-cognito-identity-provider';
import { DynamoDBClient } from '@aws-sdk/client-dynamodb';
import { DynamoDBDocumentClient, PutCommand } from '@aws-sdk/lib-dynamodb';

import type { CognitoGroupTarget, SyncStateWriter } from './clients.js';
import { normalizeEmail } from './groups.js';
import { SYNC_STATE_PK } from './constants.js';

/** Resolves emails ↔ Cognito usernames and reconciles group membership. */
export class CognitoGroupSync implements CognitoGroupTarget {
  constructor(
    private readonly userPoolId: string,
    private readonly cognito = new CognitoIdentityProviderClient({}),
  ) {}

  private emailOf(attrs: { Name?: string; Value?: string }[] | undefined): string | undefined {
    return attrs?.find((a) => a.Name === 'email')?.Value;
  }

  async ensureGroup(groupName: string): Promise<void> {
    try {
      await this.cognito.send(
        new CreateGroupCommand({ GroupName: groupName, UserPoolId: this.userPoolId }),
      );
    } catch (err) {
      // Already exists is the expected idempotent path; rethrow anything else.
      if ((err as { name?: string }).name !== 'GroupExistsException') throw err;
    }
  }

  async listMembers(groupName: string): Promise<string[]> {
    const emails: string[] = [];
    let token: string | undefined;
    do {
      const res = await this.cognito.send(
        new ListUsersInGroupCommand({
          GroupName: groupName,
          UserPoolId: this.userPoolId,
          NextToken: token,
        }),
      );
      for (const u of res.Users ?? []) {
        const email = this.emailOf(u.Attributes);
        if (email) emails.push(email);
      }
      token = res.NextToken;
    } while (token);
    return emails;
  }

  /** Resolve the Cognito username for a Workspace email (federated user). */
  private async usernameForEmail(email: string): Promise<string | undefined> {
    const res = await this.cognito.send(
      new ListUsersCommand({
        UserPoolId: this.userPoolId,
        Filter: `email = "${normalizeEmail(email)}"`,
        Limit: 1,
      }),
    );
    return res.Users?.[0]?.Username;
  }

  async addMember(groupName: string, email: string): Promise<void> {
    const username = await this.usernameForEmail(email);
    if (!username) return; // user hasn't federated into the pool yet; next sync will add them
    await this.cognito.send(
      new AdminAddUserToGroupCommand({
        GroupName: groupName,
        UserPoolId: this.userPoolId,
        Username: username,
      }),
    );
  }

  async removeMember(groupName: string, email: string): Promise<void> {
    const username = await this.usernameForEmail(email);
    if (!username) return;
    await this.cognito.send(
      new AdminRemoveUserFromGroupCommand({
        GroupName: groupName,
        UserPoolId: this.userPoolId,
        Username: username,
      }),
    );
  }
}

/** Writes the freshness marker the pre-token Lambda reads. */
export class DynamoSyncStateWriter implements SyncStateWriter {
  private readonly doc: DynamoDBDocumentClient;
  constructor(
    private readonly tableName: string,
    client: DynamoDBClient = new DynamoDBClient({}),
  ) {
    this.doc = DynamoDBDocumentClient.from(client);
  }

  async writeLastSuccessfulSync(epochMs: number): Promise<void> {
    await this.doc.send(
      new PutCommand({
        TableName: this.tableName,
        Item: {
          pk: SYNC_STATE_PK,
          lastSuccessfulSyncMs: epochMs,
          updatedAt: new Date(epochMs).toISOString(),
        },
      }),
    );
  }
}