import boto3
import json
import os
import subprocess

def awsume(profile):
    """Use awsume to switch AWS profiles."""
    result = subprocess.run(['awsume', profile], capture_output=True, text=True)
    if result.returncode != 0:
        raise Exception(f"awsume failed: {result.stderr}")
    # Parse the output to get the environment variables
    env_vars = {}
    for line in result.stdout.splitlines():
        if line.startswith('export '):
            key, value = line.split(' ')[1].split('=')
            env_vars[key] = value.strip('"')
    return env_vars

def list_s3_objects(bucket_name, prefix):
    """List all objects in an S3 bucket with the given prefix."""
    response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=prefix)
    return [obj['Key'] for obj in response.get('Contents', [])]

def restore_secret_from_s3(bucket_name, secret_name, s3_key):
    """Restore a secret from S3 to AWS Secrets Manager."""
    # Download the secret from S3
    s3_client.download_file(bucket_name, s3_key, '/tmp/secret.json')

    # Read the secret from the downloaded file
    with open('/tmp/secret.json', 'r') as file:
        secret_value = json.load(file)

    # Restore the secret to AWS Secrets Manager
    try:
        response = secrets_manager_client.create_secret(
            Name=secret_name,
            SecretString=json.dumps(secret_value)
        )
        print(f"Secret {secret_name} restored successfully.")
    except secrets_manager_client.exceptions.ResourceExistsException:
        # If the secret already exists, update it
        response = secrets_manager_client.update_secret(
            SecretId=secret_name,
            SecretString=json.dumps(secret_value)
        )
        print(f"Secret {secret_name} updated successfully.")
    except Exception as e:
        print(f"Error restoring secret {secret_name}: {e}")

if __name__ == "__main__":
    # Define the S3 bucket and prefix
    bucket_name = 'backup-orcd-secrets-manager-secrets'
    prefix = '437795906767/qa/anchore/'

    # Assume the source AWS account
    source_env_vars = awsume('source-profile')
    os.environ.update(source_env_vars)

    # Initialize boto3 clients for the source account
    s3_client = boto3.client('s3')

    # List all secrets in the specified S3 path
    s3_keys = list_s3_objects(bucket_name, prefix)

    # Assume the destination AWS account
    dest_env_vars = awsume('destination-profile')
    os.environ.update(dest_env_vars)

    # Initialize boto3 clients for the destination account
    secrets_manager_client = boto3.client('secretsmanager')

    # Restore each secret
    for s3_key in s3_keys:
        secret_name = s3_key.replace('437795906767/', '')
        restore_secret_from_s3(bucket_name, secret_name, s3_key)