thêm chức năng train trên odc predict trên planetary
This commit is contained in:
+348
@@ -0,0 +1,348 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Cognito Authentication Module
|
||||
Module xác thực sử dụng AWS Cognito Tokens
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import base64
|
||||
import boto3
|
||||
from botocore.exceptions import ClientError
|
||||
from datetime import datetime
|
||||
import jwt
|
||||
|
||||
|
||||
class CognitoAuthenticator:
|
||||
"""
|
||||
Class để xác thực và lấy AWS credentials từ Cognito tokens
|
||||
"""
|
||||
|
||||
def __init__(self, region='ap-southeast-1'):
|
||||
self.region = region
|
||||
self.cognito_identity = None
|
||||
self.access_token = None
|
||||
self.id_token = None
|
||||
self.aws_credentials = None
|
||||
|
||||
def load_tokens_from_file(self, credential_file='train_files/crediential.txt'):
|
||||
"""
|
||||
Load Cognito tokens và AWS credentials từ file
|
||||
"""
|
||||
try:
|
||||
with open(credential_file, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
aws_creds = {}
|
||||
|
||||
# Parse all credentials
|
||||
lines = content.split('\n')
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
|
||||
# Parse AWS credentials
|
||||
if line.startswith('export '):
|
||||
line = line[7:]
|
||||
if '=' in line:
|
||||
key, value = line.split('=', 1)
|
||||
value = value.strip('"')
|
||||
aws_creds[key] = value
|
||||
|
||||
# Parse Cognito tokens
|
||||
elif line.startswith('Cognito:'):
|
||||
self.access_token = line.split('Cognito:')[1].strip()
|
||||
elif line.startswith('ID:'):
|
||||
self.id_token = line.split('ID:')[1].strip()
|
||||
|
||||
# Set AWS credentials if found
|
||||
if aws_creds:
|
||||
self.aws_credentials = {
|
||||
'AccessKeyId': aws_creds.get('AWS_ACCESS_KEY_ID', ''),
|
||||
'SecretAccessKey': aws_creds.get('AWS_SECRET_ACCESS_KEY', ''),
|
||||
'SessionToken': aws_creds.get('AWS_SESSION_TOKEN', ''),
|
||||
}
|
||||
print("✓ AWS credentials loaded from file")
|
||||
|
||||
if self.access_token and self.id_token:
|
||||
print("✓ Cognito tokens loaded successfully")
|
||||
return True
|
||||
else:
|
||||
print("✗ Cognito tokens not found in file")
|
||||
return False
|
||||
|
||||
except FileNotFoundError:
|
||||
print(f"✗ Error: Credential file not found: {credential_file}")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"✗ Error loading tokens: {e}")
|
||||
return False
|
||||
|
||||
def decode_token(self, token, verify=False):
|
||||
"""
|
||||
Decode JWT token để xem thông tin
|
||||
"""
|
||||
try:
|
||||
# Decode without verification (for inspection only)
|
||||
decoded = jwt.decode(token, options={"verify_signature": False})
|
||||
return decoded
|
||||
except Exception as e:
|
||||
print(f"✗ Error decoding token: {e}")
|
||||
return None
|
||||
|
||||
def print_token_info(self):
|
||||
"""
|
||||
In thông tin từ Cognito tokens
|
||||
"""
|
||||
if not self.id_token:
|
||||
print("✗ No ID token available")
|
||||
return
|
||||
|
||||
print("\n=== Cognito Token Information ===")
|
||||
|
||||
try:
|
||||
decoded_id = self.decode_token(self.id_token)
|
||||
decoded_access = self.decode_token(self.access_token)
|
||||
|
||||
if decoded_id:
|
||||
print("\nUser Information:")
|
||||
print(f" Username: {decoded_id.get('cognito:username', 'N/A')}")
|
||||
print(f" Name: {decoded_id.get('name', 'N/A')}")
|
||||
print(f" Email: {decoded_id.get('email', 'N/A')}")
|
||||
print(f" Groups: {', '.join(decoded_id.get('cognito:groups', []))}")
|
||||
|
||||
# Check expiration
|
||||
exp = decoded_id.get('exp')
|
||||
if exp:
|
||||
exp_time = datetime.fromtimestamp(exp)
|
||||
now = datetime.now()
|
||||
if exp_time > now:
|
||||
time_left = exp_time - now
|
||||
hours = time_left.seconds // 3600
|
||||
minutes = (time_left.seconds % 3600) // 60
|
||||
print(f" Token expires: {exp_time}")
|
||||
print(f" Time remaining: {hours}h {minutes}m")
|
||||
else:
|
||||
print(f" ✗ Token EXPIRED at: {exp_time}")
|
||||
|
||||
return decoded_id, decoded_access
|
||||
|
||||
except Exception as e:
|
||||
print(f"✗ Error parsing token info: {e}")
|
||||
return None, None
|
||||
|
||||
def get_credentials_from_cognito(self, identity_pool_id=None, cognito_provider=None):
|
||||
"""
|
||||
Exchange Cognito ID token để lấy AWS temporary credentials
|
||||
|
||||
Args:
|
||||
identity_pool_id: Cognito Identity Pool ID (nếu có)
|
||||
cognito_provider: Cognito provider URL
|
||||
"""
|
||||
if not self.id_token:
|
||||
print("✗ No ID token available")
|
||||
return False
|
||||
|
||||
try:
|
||||
print("\n=== Getting AWS Credentials from Cognito ===")
|
||||
|
||||
# Nếu không có identity pool ID, thử tự động detect
|
||||
if not identity_pool_id:
|
||||
print("⚠ No Identity Pool ID provided")
|
||||
print("⚠ Using existing AWS credentials (already exchanged from Cognito)...")
|
||||
|
||||
# Kiểm tra xem có credentials đã load không
|
||||
if self.aws_credentials and self.aws_credentials.get('AccessKeyId'):
|
||||
print("✓ Using AWS credentials loaded from file")
|
||||
return True
|
||||
# Kiểm tra trong environment
|
||||
elif os.environ.get('AWS_ACCESS_KEY_ID'):
|
||||
print("✓ Using existing AWS credentials from environment")
|
||||
self.aws_credentials = {
|
||||
'AccessKeyId': os.environ.get('AWS_ACCESS_KEY_ID'),
|
||||
'SecretAccessKey': os.environ.get('AWS_SECRET_ACCESS_KEY'),
|
||||
'SessionToken': os.environ.get('AWS_SESSION_TOKEN'),
|
||||
}
|
||||
return True
|
||||
else:
|
||||
print("✗ No AWS credentials available")
|
||||
return False
|
||||
|
||||
# Create Cognito Identity client
|
||||
cognito_identity = boto3.client('cognito-identity', region_name=self.region)
|
||||
|
||||
# Default provider URL nếu không có
|
||||
if not cognito_provider:
|
||||
decoded = self.decode_token(self.id_token)
|
||||
if decoded and 'iss' in decoded:
|
||||
iss = decoded['iss']
|
||||
# Extract provider from issuer URL
|
||||
# Example: https://cognito-idp.ap-southeast-1.amazonaws.com/ap-southeast-1_C4GCbYaOa
|
||||
cognito_provider = iss.replace('https://', '')
|
||||
|
||||
print(f"Identity Pool ID: {identity_pool_id}")
|
||||
print(f"Cognito Provider: {cognito_provider}")
|
||||
|
||||
# Get identity ID
|
||||
logins = {cognito_provider: self.id_token}
|
||||
|
||||
identity_response = cognito_identity.get_id(
|
||||
IdentityPoolId=identity_pool_id,
|
||||
Logins=logins
|
||||
)
|
||||
|
||||
identity_id = identity_response['IdentityId']
|
||||
print(f"✓ Got Identity ID: {identity_id}")
|
||||
|
||||
# Get credentials for identity
|
||||
credentials_response = cognito_identity.get_credentials_for_identity(
|
||||
IdentityId=identity_id,
|
||||
Logins=logins
|
||||
)
|
||||
|
||||
self.aws_credentials = credentials_response['Credentials']
|
||||
|
||||
print("✓ Successfully obtained AWS credentials from Cognito!")
|
||||
print(f" Access Key: {self.aws_credentials['AccessKeyId'][:20]}...")
|
||||
print(f" Expiration: {self.aws_credentials['Expiration']}")
|
||||
|
||||
return True
|
||||
|
||||
except ClientError as e:
|
||||
error_code = e.response.get('Error', {}).get('Code', 'Unknown')
|
||||
print(f"✗ AWS Error: {error_code}")
|
||||
print(f"✗ Message: {e.response.get('Error', {}).get('Message', 'Unknown')}")
|
||||
|
||||
if error_code == 'NotAuthorizedException':
|
||||
print("✗ Token may be expired or invalid")
|
||||
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"✗ Error getting credentials: {e}")
|
||||
return False
|
||||
|
||||
def set_environment_credentials(self):
|
||||
"""
|
||||
Set AWS credentials vào environment variables
|
||||
"""
|
||||
if not self.aws_credentials:
|
||||
print("✗ No AWS credentials available")
|
||||
return False
|
||||
|
||||
try:
|
||||
os.environ['AWS_ACCESS_KEY_ID'] = self.aws_credentials['AccessKeyId']
|
||||
os.environ['AWS_SECRET_ACCESS_KEY'] = self.aws_credentials['SecretAccessKey']
|
||||
os.environ['AWS_SESSION_TOKEN'] = self.aws_credentials['SessionToken']
|
||||
|
||||
print("✓ AWS credentials set in environment")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"✗ Error setting credentials: {e}")
|
||||
return False
|
||||
|
||||
def test_s3_access(self, bucket_name='sentinel-cogs', region='us-west-2'):
|
||||
"""
|
||||
Test S3 access với credentials từ Cognito
|
||||
"""
|
||||
try:
|
||||
print(f"\n=== Testing S3 Access ===")
|
||||
print(f"Bucket: {bucket_name}")
|
||||
|
||||
# Create S3 client với credentials
|
||||
if self.aws_credentials:
|
||||
s3_client = boto3.client(
|
||||
's3',
|
||||
region_name=region,
|
||||
aws_access_key_id=self.aws_credentials['AccessKeyId'],
|
||||
aws_secret_access_key=self.aws_credentials['SecretAccessKey'],
|
||||
aws_session_token=self.aws_credentials['SessionToken']
|
||||
)
|
||||
else:
|
||||
# Use credentials from environment
|
||||
s3_client = boto3.client('s3', region_name=region)
|
||||
|
||||
# Test bucket access
|
||||
s3_client.head_bucket(Bucket=bucket_name)
|
||||
print(f"✓ Successfully accessed bucket: {bucket_name}")
|
||||
|
||||
# List some objects
|
||||
response = s3_client.list_objects_v2(
|
||||
Bucket=bucket_name,
|
||||
MaxKeys=100
|
||||
)
|
||||
|
||||
if 'Contents' in response:
|
||||
print(f"✓ Listed {len(response['Contents'])} objects:")
|
||||
for obj in response['Contents']:
|
||||
size_mb = obj['Size'] / (1024 * 1024)
|
||||
print(f" - {obj['Key']} ({size_mb:.2f} MB)")
|
||||
|
||||
return True
|
||||
|
||||
except ClientError as e:
|
||||
error_code = e.response.get('Error', {}).get('Code', 'Unknown')
|
||||
print(f"✗ S3 Error: {error_code}")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"✗ Error: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
Test function
|
||||
"""
|
||||
print("=" * 60)
|
||||
print("Cognito Authentication Test")
|
||||
print("Test Xác Thực Cognito")
|
||||
print("=" * 60)
|
||||
|
||||
# Initialize authenticator
|
||||
auth = CognitoAuthenticator(region='ap-southeast-1')
|
||||
|
||||
# Step 1: Load tokens
|
||||
print("\n[Step 1] Loading Cognito tokens...")
|
||||
if not auth.load_tokens_from_file():
|
||||
print("\n✗ Failed to load tokens")
|
||||
return False
|
||||
|
||||
# Step 2: Display token info
|
||||
print("\n[Step 2] Parsing token information...")
|
||||
auth.print_token_info()
|
||||
|
||||
# Step 3: Get AWS credentials
|
||||
print("\n[Step 3] Getting AWS credentials...")
|
||||
|
||||
# Option A: Nếu có Identity Pool ID (uncomment nếu biết)
|
||||
# identity_pool_id = 'ap-southeast-1:xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx'
|
||||
# auth.get_credentials_from_cognito(identity_pool_id=identity_pool_id)
|
||||
|
||||
# Option B: Sử dụng credentials có sẵn trong file
|
||||
auth.get_credentials_from_cognito()
|
||||
|
||||
# Step 4: Set environment
|
||||
print("\n[Step 4] Setting environment credentials...")
|
||||
auth.set_environment_credentials()
|
||||
|
||||
# Step 5: Test S3 access
|
||||
print("\n[Step 5] Testing S3 access...")
|
||||
auth.test_s3_access()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("✓ Test completed!")
|
||||
print("=" * 60)
|
||||
|
||||
return True
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
success = main()
|
||||
exit(0 if success else 1)
|
||||
except KeyboardInterrupt:
|
||||
print("\n\n✗ Test interrupted by user")
|
||||
exit(1)
|
||||
except Exception as e:
|
||||
print(f"\n✗ Fatal error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
exit(1)
|
||||
Reference in New Issue
Block a user