Merge pull request #2808 from trycompai/feat/support-gov-cloud-aws

[dev] [tofikwest] feat/support-gov-cloud-aws

Author: tofikwest

apps/api/.env.example

@@ -53,4 +53,9 @@ BACKGROUND_CHECK_WEBHOOK_SECRET=
 BACKGROUND_WH_ENDPOINT=
 STRIPE_BACKGROUND_CHECK_PRICE_ID=price_1TRWckCkFWhKYvHIA1GLv1sO
 
+# We are not using security hub service it is only the variable prefix for matching with old one
 SECURITY_HUB_ROLE_ASSUMER_ARN=
+SECURITY_HUB_GOVCLOUD_ROLE_ASSUMER_ARN=
+SECURITY_HUB_GOVCLOUD_ACCESS_KEY_ID=
+SECURITY_HUB_GOVCLOUD_SECRET_ACCESS_KEY=
+SECURITY_HUB_GOVCLOUD_SESSION_TOKEN=

apps/api/src/cloud-security/ai-remediation.prompt.ts

@@ -99,6 +99,28 @@ export type CompletePermissions = z.infer<typeof completePermissionsSchema>;
 
 // ─── Prompt Builders ────────────────────────────────────────────────────────
 
+function inferAwsPartition(finding: {
+  resourceId: string;
+  evidence: Record<string, unknown>;
+}): 'aws' | 'aws-us-gov' {
+  if (finding.resourceId.startsWith('arn:aws-us-gov:')) return 'aws-us-gov';
+
+  const region =
+    typeof finding.evidence.region === 'string' ? finding.evidence.region : '';
+  return region.startsWith('us-gov-') ? 'aws-us-gov' : 'aws';
+}
+
+function inferAwsRegion(finding: {
+  resourceId: string;
+  evidence: Record<string, unknown>;
+}): string {
+  if (typeof finding.evidence.region === 'string')
+    return finding.evidence.region;
+
+  const arnMatch = finding.resourceId.match(/^arn:[^:]+:[^:]*:([^:]*):/);
+  return arnMatch?.[1] || 'the execution region';
+}
+
 const SYSTEM_PROMPT = `You are an AWS security remediation expert. You analyze security findings and produce structured fix plans that will be executed by an automated system using AWS SDK v3.
 
 A human will ALWAYS review your plan before execution. Be precise and correct.
@@ -117,12 +139,20 @@ A human will ALWAYS review your plan before execution. Be precise and correct.
 
 ## RESOURCE ID PARSING
 - Extract actual resource names from ARNs:
-  - "arn:aws:s3:::my-bucket" → Bucket: "my-bucket"
-  - "arn:aws:kms:us-east-1:123:key/abc" → KeyId: "arn:aws:kms:us-east-1:123:key/abc" (use full ARN for KMS)
-  - "arn:aws:rds:us-east-1:123:db:mydb" → DBInstanceIdentifier: "mydb"
-  - "arn:aws:ec2:us-east-1:123:vpc/vpc-abc" → VpcId: "vpc-abc"
+  - "arn:aws:s3:::my-bucket" or "arn:aws-us-gov:s3:::my-bucket" → Bucket: "my-bucket"
+  - "arn:aws:kms:us-east-1:123:key/abc" → KeyId: use the full ARN exactly as provided
+  - "arn:aws-us-gov:kms:us-gov-west-1:123:key/abc" → KeyId: use the full GovCloud ARN exactly as provided
+  - "arn:aws:rds:us-east-1:123:db:mydb" or "arn:aws-us-gov:rds:us-gov-west-1:123:db:mydb" → DBInstanceIdentifier: "mydb"
+  - "arn:aws:ec2:us-east-1:123:vpc/vpc-abc" or "arn:aws-us-gov:ec2:us-gov-west-1:123:vpc/vpc-abc" → VpcId: "vpc-abc"
 - Use the correct parameter names that the AWS SDK expects
 
+## AWS PARTITIONS AND GOVCLOUD
+- Preserve the AWS partition from the finding context.
+- If AWS Partition is "aws-us-gov", every ARN you create or pass MUST start with "arn:aws-us-gov:".
+- If AWS Partition is "aws", every ARN you create or pass MUST start with "arn:aws:".
+- Never convert a GovCloud ARN to a commercial AWS ARN.
+- For GovCloud findings, use GovCloud regions such as "us-gov-west-1" or "us-gov-east-1"; never default to "us-east-1".
+
 ## SAFETY RULES (NEVER violate)
 - NEVER delete data, buckets, tables, databases, or file systems
 - NEVER modify IAM policies, roles, or users in ways that could lock out users
@@ -259,10 +289,19 @@ export function buildFixPlanPrompt(finding: {
   findingKey: string;
   evidence: Record<string, unknown>;
 }): string {
+  const awsPartition = inferAwsPartition(finding);
+  const awsRegion = inferAwsRegion(finding);
+
   return `Analyze this AWS security finding and generate a fix plan.
 
 IMPORTANT: Your fix must change the EXACT AWS setting/resource that caused this finding. The scan will re-check the same thing after the fix — if you fix something different, the finding will persist.
 
+AWS EXECUTION CONTEXT:
+- AWS Partition: ${awsPartition}
+- Region: ${awsRegion}
+- When constructing ARNs, use partition prefix: arn:${awsPartition}:
+- If region-specific values are needed, use this region unless the finding explicitly gives a different one.
+
 FINDING:
 - Title: ${finding.title}
 - Description: ${finding.description ?? 'N/A'}

apps/api/src/cloud-security/aws-command-executor.ts

@@ -1,5 +1,9 @@
 import type { AwsCredentialIdentity } from '@aws-sdk/types';
 import type { AwsCommandStep } from './ai-remediation.prompt';
+import {
+  type AwsPartition,
+  getAwsPartitionForRegion,
+} from './aws-partition.utils';
 
 import * as s3 from '@aws-sdk/client-s3';
 import * as dynamodb from '@aws-sdk/client-dynamodb';
@@ -145,6 +149,50 @@ const JSON_STRING_PARAMS = new Set([
   'Definition',
 ]);
 
+function normalizeArnPartition(value: string, partition: AwsPartition): string {
+  if (partition === 'aws-us-gov') {
+    return value.replace(/\barn:aws:/g, 'arn:aws-us-gov:');
+  }
+
+  return value;
+}
+
+function normalizeArnPartitionsInValue(
+  value: unknown,
+  partition: AwsPartition,
+): unknown {
+  if (typeof value === 'string') {
+    return normalizeArnPartition(value, partition);
+  }
+
+  if (Array.isArray(value)) {
+    return value.map((item) => normalizeArnPartitionsInValue(item, partition));
+  }
+
+  if (value !== null && typeof value === 'object') {
+    return Object.fromEntries(
+      Object.entries(value).map(([key, item]) => [
+        key,
+        normalizeArnPartitionsInValue(item, partition),
+      ]),
+    );
+  }
+
+  return value;
+}
+
+function normalizeArnPartitions(
+  input: Record<string, unknown>,
+  region: string,
+): void {
+  const partition = getAwsPartitionForRegion(region);
+  if (partition === 'aws') return;
+
+  for (const [key, value] of Object.entries(input)) {
+    input[key] = normalizeArnPartitionsInValue(value, partition);
+  }
+}
+
 /**
  * Universal pre-execution param normalisation.
  * Fixes common AI mistakes without per-command logic.
@@ -154,6 +202,8 @@ function normaliseInputParams(
   command: string,
   region: string,
 ): void {
+  normalizeArnPartitions(input, region);
+
   for (const [key, value] of Object.entries(input)) {
     // Rule 1: Stringify any object param that AWS expects as a JSON string
     if (

apps/api/src/cloud-security/aws-partition.utils.spec.ts

@@ -0,0 +1,73 @@
+import {
+  getAwsBaseCredentials,
+  getAwsDefaultRegion,
+  getAwsPartitionForRegion,
+  parseAwsRoleArn,
+  validateAwsPartitionConfig,
+} from './aws-partition.utils';
+
+describe('aws partition utils', () => {
+  it('uses GovCloud defaults for the aws-us-gov partition', () => {
+    expect(getAwsDefaultRegion('aws')).toBe('us-east-1');
+    expect(getAwsDefaultRegion('aws-us-gov')).toBe('us-gov-west-1');
+    expect(getAwsPartitionForRegion('us-east-1')).toBe('aws');
+    expect(getAwsPartitionForRegion('us-gov-east-1')).toBe('aws-us-gov');
+  });
+
+  it('parses commercial and GovCloud role ARNs', () => {
+    expect(
+      parseAwsRoleArn('arn:aws:iam::123456789012:role/CompAI-Auditor'),
+    ).toEqual({ partition: 'aws', accountId: '123456789012' });
+    expect(
+      parseAwsRoleArn('arn:aws-us-gov:iam::123456789012:role/CompAI-Auditor'),
+    ).toEqual({ partition: 'aws-us-gov', accountId: '123456789012' });
+  });
+
+  it('rejects mismatched role ARN and region partitions', () => {
+    expect(
+      validateAwsPartitionConfig({
+        partition: 'aws-us-gov',
+        roleArn: 'arn:aws:iam::123456789012:role/CompAI-Auditor',
+        regions: ['us-gov-west-1', 'us-east-1'],
+      }),
+    ).toEqual([
+      'IAM Role ARN partition (aws) must match selected AWS environment (aws-us-gov).',
+      'Selected regions do not match aws-us-gov: us-east-1.',
+    ]);
+  });
+
+  it('accepts commercial and GovCloud configurations independently', () => {
+    expect(
+      validateAwsPartitionConfig({
+        partition: 'aws',
+        roleArn: 'arn:aws:iam::123456789012:role/CompAI-Auditor',
+        regions: ['us-east-1', 'us-west-2'],
+      }),
+    ).toEqual([]);
+
+    expect(
+      validateAwsPartitionConfig({
+        partition: 'aws-us-gov',
+        roleArn: 'arn:aws-us-gov:iam::123456789012:role/CompAI-Auditor',
+        regions: ['us-gov-west-1', 'us-gov-east-1'],
+      }),
+    ).toEqual([]);
+  });
+
+  it('uses explicit GovCloud base credentials when configured', () => {
+    process.env.SECURITY_HUB_GOVCLOUD_ACCESS_KEY_ID = 'AKIAGOV';
+    process.env.SECURITY_HUB_GOVCLOUD_SECRET_ACCESS_KEY = 'secret';
+    process.env.SECURITY_HUB_GOVCLOUD_SESSION_TOKEN = 'token';
+
+    expect(getAwsBaseCredentials('aws-us-gov')).toEqual({
+      accessKeyId: 'AKIAGOV',
+      secretAccessKey: 'secret',
+      sessionToken: 'token',
+    });
+    expect(getAwsBaseCredentials('aws')).toBeUndefined();
+
+    delete process.env.SECURITY_HUB_GOVCLOUD_ACCESS_KEY_ID;
+    delete process.env.SECURITY_HUB_GOVCLOUD_SECRET_ACCESS_KEY;
+    delete process.env.SECURITY_HUB_GOVCLOUD_SESSION_TOKEN;
+  });
+});

apps/api/src/cloud-security/aws-partition.utils.ts

@@ -0,0 +1,100 @@
+import type { AwsCredentialIdentity } from '@aws-sdk/types';
+
+export type AwsPartition = 'aws' | 'aws-us-gov';
+
+const AWS_PARTITIONS = new Set<AwsPartition>(['aws', 'aws-us-gov']);
+
+export function normalizeAwsPartition(value: unknown): AwsPartition {
+  return typeof value === 'string' && AWS_PARTITIONS.has(value as AwsPartition)
+    ? (value as AwsPartition)
+    : 'aws';
+}
+
+export function getAwsDefaultRegion(partition: AwsPartition): string {
+  return partition === 'aws-us-gov' ? 'us-gov-west-1' : 'us-east-1';
+}
+
+export function getAwsPartitionForRegion(region: string): AwsPartition {
+  return region.startsWith('us-gov-') ? 'aws-us-gov' : 'aws';
+}
+
+export function getAwsRoleAssumerEnvName(partition: AwsPartition): string {
+  return partition === 'aws-us-gov'
+    ? 'SECURITY_HUB_GOVCLOUD_ROLE_ASSUMER_ARN'
+    : 'SECURITY_HUB_ROLE_ASSUMER_ARN';
+}
+
+export function getAwsRoleAssumerArn(partition: AwsPartition): string | undefined {
+  return process.env[getAwsRoleAssumerEnvName(partition)];
+}
+
+export function getAwsBaseCredentials(
+  partition: AwsPartition,
+): AwsCredentialIdentity | undefined {
+  if (partition !== 'aws-us-gov') return undefined;
+
+  const accessKeyId = process.env.SECURITY_HUB_GOVCLOUD_ACCESS_KEY_ID;
+  const secretAccessKey =
+    process.env.SECURITY_HUB_GOVCLOUD_SECRET_ACCESS_KEY;
+  if (!accessKeyId || !secretAccessKey) return undefined;
+
+  return {
+    accessKeyId,
+    secretAccessKey,
+    sessionToken: process.env.SECURITY_HUB_GOVCLOUD_SESSION_TOKEN,
+  };
+}
+
+export function parseAwsRoleArn(
+  roleArn: string,
+): { partition: AwsPartition; accountId: string } | null {
+  const match = roleArn.match(/^arn:(aws|aws-us-gov):iam::(\d{12}):role\/.+$/);
+  if (!match) return null;
+
+  return {
+    partition: match[1] as AwsPartition,
+    accountId: match[2],
+  };
+}
+
+export function validateAwsPartitionConfig(params: {
+  partition: AwsPartition;
+  roleArn: string;
+  regions: string[];
+  remediationRoleArn?: string;
+}): string[] {
+  const errors: string[] = [];
+  const parsedRoleArn = parseAwsRoleArn(params.roleArn);
+
+  if (!parsedRoleArn) {
+    errors.push(
+      'Invalid IAM Role ARN format. Expected: arn:aws:iam::ACCOUNT_ID:role/ROLE_NAME or arn:aws-us-gov:iam::ACCOUNT_ID:role/ROLE_NAME',
+    );
+  } else if (parsedRoleArn.partition !== params.partition) {
+    errors.push(
+      `IAM Role ARN partition (${parsedRoleArn.partition}) must match selected AWS environment (${params.partition}).`,
+    );
+  }
+
+  if (params.remediationRoleArn) {
+    const parsedRemediationArn = parseAwsRoleArn(params.remediationRoleArn);
+    if (!parsedRemediationArn) {
+      errors.push('Invalid Remediation Role ARN format.');
+    } else if (parsedRemediationArn.partition !== params.partition) {
+      errors.push(
+        `Remediation Role ARN partition (${parsedRemediationArn.partition}) must match selected AWS environment (${params.partition}).`,
+      );
+    }
+  }
+
+  const mismatchedRegions = params.regions.filter(
+    (region) => getAwsPartitionForRegion(region) !== params.partition,
+  );
+  if (mismatchedRegions.length > 0) {
+    errors.push(
+      `Selected regions do not match ${params.partition}: ${mismatchedRegions.join(', ')}.`,
+    );
+  }
+
+  return errors;
+}

apps/api/src/cloud-security/cloud-security-query.service.ts

@@ -31,6 +31,7 @@ export interface CloudProvider {
   variables: Record<string, unknown> | null;
   requiredVariables: string[];
   accountId?: string;
+  awsType?: string;
   regions?: string[];
   tenantId?: string;
   subscriptionId?: string;
@@ -133,6 +134,8 @@ export class CloudSecurityQueryService {
           typeof metadata.accountId === 'string'
             ? metadata.accountId
             : undefined,
+        awsType:
+          typeof metadata.awsType === 'string' ? metadata.awsType : undefined,
         regions: Array.isArray(metadata.regions)
           ? metadata.regions.filter((r): r is string => typeof r === 'string')
           : undefined,
@@ -171,6 +174,8 @@ export class CloudSecurityQueryService {
           typeof settings.accountId === 'string'
             ? settings.accountId
             : undefined,
+        awsType:
+          typeof settings.awsType === 'string' ? settings.awsType : undefined,
         regions: Array.isArray(settings.regions)
           ? settings.regions.filter((r): r is string => typeof r === 'string')
           : undefined,