Add: 添加Session Token支持

backend/app/api/credentials.py

@@ -427,6 +427,10 @@ def validate_credential():
                 'access_key_id': base_config.access_key_id,
                 'secret_access_key': base_config.get_secret_access_key()
             }
+            # Add session token if available
+            session_token = base_config.get_session_token()
+            if session_token:
+                base_credentials['session_token'] = session_token
         else:
             credential_config = {
                 'access_key_id': credential.access_key_id,
@@ -463,6 +467,10 @@ def validate_credential():
                 'access_key_id': base_config.access_key_id,
                 'secret_access_key': base_config.get_secret_access_key()
             }
+            # Add session token if available
+            session_token = base_config.get_session_token()
+            if session_token:
+                base_credentials['session_token'] = session_token
         elif credential_type == 'access_key':
             access_key_id = data.get('access_key_id', '').strip()
             secret_access_key = data.get('secret_access_key', '').strip()
@@ -542,7 +550,8 @@ def update_base_role():
     Request body:
         {
             "access_key_id": "string" (required),
-            "secret_access_key": "string" (required)
+            "secret_access_key": "string" (required),
+            "session_token": "string" (optional, for temporary credentials)
         }
     
     Returns:
@@ -559,6 +568,7 @@ def update_base_role():
     # Validate required fields
     access_key_id = data.get('access_key_id', '').strip()
     secret_access_key = data.get('secret_access_key', '').strip()
+    session_token = data.get('session_token', '').strip() if data.get('session_token') else None
     
     if not access_key_id:
         raise ValidationError(
@@ -574,12 +584,16 @@ def update_base_role():
     
     # Validate the credentials before saving
     try:
+        credential_config = {
+            'access_key_id': access_key_id,
+            'secret_access_key': secret_access_key
+        }
+        if session_token:
+            credential_config['session_token'] = session_token
+        
         provider = AWSCredentialProvider(
             credential_type='access_key',
-            credential_config={
-                'access_key_id': access_key_id,
-                'secret_access_key': secret_access_key
-            }
+            credential_config=credential_config
         )
         provider.validate()
     except CredentialError as e:
@@ -600,12 +614,14 @@ def update_base_role():
         # Update existing config
         config.access_key_id = access_key_id
         config.set_secret_access_key(secret_access_key)
+        config.set_session_token(session_token)
     else:
         # Create new config
         config = BaseAssumeRoleConfig(
             access_key_id=access_key_id
         )
         config.set_secret_access_key(secret_access_key)
+        config.set_session_token(session_token)
         db.session.add(config)
     
     db.session.commit()

backend/app/models/credential.py

@@ -99,6 +99,7 @@ class BaseAssumeRoleConfig(db.Model):
     id = db.Column(db.Integer, primary_key=True)
     access_key_id = db.Column(db.String(255), nullable=False)
     secret_access_key_encrypted = db.Column(db.Text, nullable=False)
+    session_token_encrypted = db.Column(db.Text, nullable=True)  # Optional session token for temporary credentials
     updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
     
     def set_secret_access_key(self, secret_key: str) -> None:
@@ -111,11 +112,25 @@ class BaseAssumeRoleConfig(db.Model):
             return decrypt_value(self.secret_access_key_encrypted)
         return None
     
+    def set_session_token(self, session_token: str) -> None:
+        """Encrypt and store the session token"""
+        if session_token:
+            self.session_token_encrypted = encrypt_value(session_token)
+        else:
+            self.session_token_encrypted = None
+    
+    def get_session_token(self) -> str:
+        """Decrypt and return the session token"""
+        if self.session_token_encrypted:
+            return decrypt_value(self.session_token_encrypted)
+        return None
+    
     def to_dict(self, mask_sensitive: bool = True) -> dict:
         """Convert config to dictionary"""
         result = {
             'id': self.id,
             'access_key_id': self.access_key_id,
+            'has_session_token': self.session_token_encrypted is not None,
             'updated_at': format_datetime(self.updated_at)
         }
         

backend/app/scanners/credentials.py

@@ -78,6 +78,7 @@ class AWSCredentialProvider:
         Get a session using Assume Role.
         
         Uses the base account credentials to assume a role in the target account.
+        Supports optional session token for temporary base credentials.
         """
         role_arn = self.credential_config.get('role_arn')
         external_id = self.credential_config.get('external_id')
@@ -90,11 +91,16 @@ class AWSCredentialProvider:
         
         try:
             # Create base session with the centralized account credentials
-            base_session = boto3.Session(
-                aws_access_key_id=self.base_credentials.get('access_key_id'),
-                aws_secret_access_key=self.base_credentials.get('secret_access_key'),
-                region_name=region_name or 'us-east-1'
-            )
+            base_session_params = {
+                'aws_access_key_id': self.base_credentials.get('access_key_id'),
+                'aws_secret_access_key': self.base_credentials.get('secret_access_key'),
+                'region_name': region_name or 'us-east-1'
+            }
+            # Add session token if provided (for temporary credentials)
+            if self.base_credentials.get('session_token'):
+                base_session_params['aws_session_token'] = self.base_credentials.get('session_token')
+            
+            base_session = boto3.Session(**base_session_params)
             
             # Use STS to assume the role
             sts_client = base_session.client('sts')
@@ -133,19 +139,26 @@ class AWSCredentialProvider:
     def _get_access_key_session(self, region_name: Optional[str] = None) -> boto3.Session:
         """
         Get a session using Access Key credentials.
+        Supports optional session token for temporary credentials.
         """
         access_key_id = self.credential_config.get('access_key_id')
         secret_access_key = self.credential_config.get('secret_access_key')
+        session_token = self.credential_config.get('session_token')
         
         if not access_key_id or not secret_access_key:
             raise CredentialError("Access Key ID and Secret Access Key are required")
         
         try:
-            return boto3.Session(
-                aws_access_key_id=access_key_id,
-                aws_secret_access_key=secret_access_key,
-                region_name=region_name
-            )
+            session_params = {
+                'aws_access_key_id': access_key_id,
+                'aws_secret_access_key': secret_access_key,
+                'region_name': region_name
+            }
+            # Add session token if provided (for temporary credentials)
+            if session_token:
+                session_params['aws_session_token'] = session_token
+            
+            return boto3.Session(**session_params)
         except Exception as e:
             logger.error(f"Failed to create session with access key: {str(e)}")
             raise CredentialError(f"Failed to create session: {str(e)}")
@@ -222,6 +235,10 @@ def create_credential_provider_from_model(
                 'access_key_id': base_config.access_key_id,
                 'secret_access_key': base_config.get_secret_access_key()
             }
+            # Add session token if available (for temporary credentials)
+            session_token = base_config.get_session_token()
+            if session_token:
+                base_credentials['session_token'] = session_token
     else:  # access_key
         credential_config = {
             'access_key_id': credential.access_key_id,

backend/migrations/versions/003_add_session_token_to_base_role.py

@@ -0,0 +1,27 @@
+"""Add session_token to base_assume_role_config
+
+Revision ID: 003_add_session_token
+Revises: 002_make_account_id_nullable
+Create Date: 2026-01-23
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '003_add_session_token'
+down_revision = '002_make_account_id_nullable'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # Add session_token_encrypted column to base_assume_role_config table
+    op.add_column('base_assume_role_config', 
+                  sa.Column('session_token_encrypted', sa.Text(), nullable=True))
+
+
+def downgrade():
+    # Remove session_token_encrypted column
+    op.drop_column('base_assume_role_config', 'session_token_encrypted')

frontend/src/pages/Credentials.tsx

@@ -266,6 +266,7 @@ export default function Credentials() {
       await credentialService.updateBaseRole({
         accessKeyId: values.accessKeyId,
         secretAccessKey: values.secretAccessKey,
+        sessionToken: values.sessionToken || '',
       });
       message.success('Base Assume Role configuration saved successfully');
       setBaseRoleModalVisible(false);
@@ -663,7 +664,7 @@ export default function Credentials() {
       >
         <Alert
           message="About Base Assume Role"
-          description="This is the base AWS account credentials used to assume roles in target accounts. All Assume Role type credentials will use these base credentials to perform STS AssumeRole operations."
+          description="This is the base AWS account credentials used to assume roles in target accounts. All Assume Role type credentials will use these base credentials to perform STS AssumeRole operations. You can optionally provide a Session Token if using temporary credentials."
           type="info"
           showIcon
           style={{ marginBottom: 16 }}
@@ -675,6 +676,13 @@ export default function Credentials() {
               <Descriptions.Item label="Current Access Key ID">
                 {baseRoleConfig.accessKeyId || 'Not configured'}
               </Descriptions.Item>
+              <Descriptions.Item label="Session Token">
+                {baseRoleConfig.hasSessionToken ? (
+                  <Tag color="blue">Configured</Tag>
+                ) : (
+                  <Tag>Not configured</Tag>
+                )}
+              </Descriptions.Item>
               <Descriptions.Item label="Status">
                 <Tag color="success">Configured</Tag>
               </Descriptions.Item>
@@ -705,6 +713,14 @@ export default function Credentials() {
             <Input.Password placeholder="Enter Secret Access Key" />
           </Form.Item>
 
+          <Form.Item
+            name="sessionToken"
+            label="Session Token"
+            help="Optional. Required only for temporary credentials (e.g., from STS AssumeRole or GetSessionToken)"
+          >
+            <Input.Password placeholder="Enter Session Token (optional)" />
+          </Form.Item>
+
           <Form.Item style={{ marginBottom: 0, textAlign: 'right' }}>
             <Space>
               <Button onClick={() => setBaseRoleModalVisible(false)}>Cancel</Button>

frontend/src/services/credentials.ts

@@ -14,6 +14,8 @@ export interface CreateCredentialRequest {
 export interface BaseRoleConfig {
   accessKeyId: string;
   secretAccessKey?: string;
+  sessionToken?: string;
+  hasSessionToken?: boolean;
 }
 
 export interface BaseRoleResponse {
@@ -88,9 +90,14 @@ export const credentialService = {
   },
 
   updateBaseRole: async (data: BaseRoleConfig): Promise<void> => {
-    await api.post('/credentials/base-role', {
+    const payload: Record<string, string> = {
       access_key_id: data.accessKeyId,
-      secret_access_key: data.secretAccessKey,
-    });
+      secret_access_key: data.secretAccessKey!,
+    };
+    // Include session token if provided (can be empty string to clear it)
+    if (data.sessionToken !== undefined) {
+      payload.session_token = data.sessionToken;
+    }
+    await api.post('/credentials/base-role', payload);
   },
 };