// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.components.service
import com.intellij.openapi.extensions.ExtensionPointName
import com.intellij.openapi.util.SimpleModificationTracker
import com.intellij.util.messages.MessageBus
import com.intellij.util.messages.Topic
import software.amazon.awssdk.auth.credentials.AwsCredentials
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider
import software.aws.toolkits.core.credentials.CredentialIdentifier
import software.aws.toolkits.core.credentials.CredentialProviderFactory
import software.aws.toolkits.core.credentials.CredentialProviderNotFoundException
import software.aws.toolkits.core.credentials.ToolkitCredentialsChangeListener
import software.aws.toolkits.core.credentials.ToolkitCredentialsProvider
import software.aws.toolkits.core.region.AwsRegion
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.tryOrNull
import software.aws.toolkits.jetbrains.utils.runUnderProgressIfNeeded
import software.aws.toolkits.resources.message
import software.aws.toolkits.telemetry.AwsTelemetry
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.atomic.AtomicInteger

abstract class CredentialManager : SimpleModificationTracker() {
    private val providerIds = ConcurrentHashMap<String, CredentialIdentifier>()
    private val awsCredentialProviderCache = ConcurrentHashMap<String, ConcurrentHashMap<String, AwsCredentialsProvider>>()

    protected abstract fun factoryMapping(): Map<String, CredentialProviderFactory>

    @Throws(CredentialProviderNotFoundException::class)
    fun getAwsCredentialProvider(providerId: CredentialIdentifier, region: AwsRegion): ToolkitCredentialsProvider =
        ToolkitCredentialsProvider(providerId, AwsCredentialProviderProxy(providerId.id, region))

    fun getCredentialIdentifiers(): List<CredentialIdentifier> = providerIds.values
        .sortedWith(compareBy(String.CASE_INSENSITIVE_ORDER) { it.displayName })

    fun getCredentialIdentifierById(id: String): CredentialIdentifier? = providerIds[id]

    // TODO: Convert these to bulk listeners so we only send N messages where N is # of extensions vs # of providers
    protected fun addProvider(identifier: CredentialIdentifier) {
        providerIds[identifier.id] = identifier

        incModificationCount()
        ApplicationManager.getApplication().messageBus.syncPublisher(CREDENTIALS_CHANGED).providerAdded(identifier)
    }

    protected fun modifyProvider(identifier: CredentialIdentifier) {
        awsCredentialProviderCache.remove(identifier.id)
        providerIds[identifier.id] = identifier

        incModificationCount()
        ApplicationManager.getApplication().messageBus.syncPublisher(CREDENTIALS_CHANGED).providerModified(identifier)
    }

    protected fun removeProvider(identifier: CredentialIdentifier) {
        providerIds.remove(identifier.id)
        awsCredentialProviderCache.remove(identifier.id)

        incModificationCount()
        ApplicationManager.getApplication().messageBus.syncPublisher(CREDENTIALS_CHANGED).providerRemoved(identifier)
    }

    /**
     * Inner class that lazy-ily requests the true AwsCredentialsProvider from the factory when needed. This acts as a middle man that allows for existing
     * ToolkitCredentialsProvider (like ones passed to existing SDK clients), to keep operating even if the credentials they represent have been updated such
     * as loading from disk when new values.
     */
    private inner class AwsCredentialProviderProxy(private val providerId: String, private val region: AwsRegion) : AwsCredentialsProvider {
        override fun resolveCredentials(): AwsCredentials = runUnderProgressIfNeeded(null, message("credentials.retrieving"), cancelable = true) {
            getOrCreateAwsCredentialsProvider(providerId, region).resolveCredentials()
        }

        private fun getOrCreateAwsCredentialsProvider(providerId: String, region: AwsRegion): AwsCredentialsProvider {
            // Validate that the provider ID is still valid and get the latest copy
            val identifier = providerIds[providerId]
                ?: throw CredentialProviderNotFoundException("Provider ID $providerId was removed, can't resolve credentials")

            val partitionCache = awsCredentialProviderCache.computeIfAbsent(providerId) { ConcurrentHashMap() }

            // If we already resolved creds for this partition and provider ID, just return it, else compute new one
            return partitionCache.computeIfAbsent(region.partitionId) {
                val providerFactory = factoryMapping()[identifier.factoryId]
                    ?: throw CredentialProviderNotFoundException("No provider factory found with ID ${identifier.factoryId}")

                try {
                    providerFactory.createAwsCredentialProvider(identifier, region)
                } catch (e: Exception) {
                    throw CredentialProviderNotFoundException("Failed to create underlying AwsCredentialProvider", e)
                }
            }
        }
    }

    companion object {
        @JvmStatic
        fun getInstance(): CredentialManager = service()

        /***
         * [MessageBus] topic for when credential providers get added/changed/deleted
         */
        val CREDENTIALS_CHANGED: Topic<ToolkitCredentialsChangeListener> = Topic.create(
            "AWS toolkit credential providers changed",
            ToolkitCredentialsChangeListener::class.java
        )
    }
}

class DefaultCredentialManager : CredentialManager() {
    private val extensionMap: Map<String, CredentialProviderFactory>
        get() = EP_NAME.extensionList.associateBy {
            it.id
        }

    init {
        extensionMap.values.forEach { providerFactory ->
            val count = AtomicInteger(0)
            LOG.tryOrNull("Failed to set up $providerFactory") {
                providerFactory.setUp { change ->
                    change.added.forEach {
                        addProvider(it)
                        count.incrementAndGet()
                    }

                    change.modified.forEach {
                        modifyProvider(it)
                    }

                    change.removed.forEach {
                        removeProvider(it)
                        count.decrementAndGet()
                    }
                    AwsTelemetry.loadCredentials(
                        credentialSourceId = providerFactory.credentialSourceId.toTelemetryCredentialSourceId(),
                        value = count.get().toDouble()
                    )
                }
            }
        }
    }

    override fun factoryMapping(): Map<String, CredentialProviderFactory> = extensionMap

    companion object {
        val EP_NAME = ExtensionPointName.create<CredentialProviderFactory>("aws.toolkit.credentialProviderFactory")
        private val LOG = getLogger<DefaultCredentialManager>()
    }
}


// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.openapi.fileEditor.FileDocumentManager
import com.intellij.openapi.fileEditor.FileEditor
import com.intellij.openapi.help.HelpManager
import com.intellij.openapi.project.DumbAware
import com.intellij.openapi.project.Project
import com.intellij.openapi.util.Key
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.ui.EditorNotificationPanel
import com.intellij.ui.EditorNotifications
import software.amazon.awssdk.profiles.ProfileFileLocation
import software.aws.toolkits.jetbrains.core.credentials.CredentialsFileHelpNotificationProvider.CredentialFileNotificationPanel
import software.aws.toolkits.jetbrains.core.help.HelpIds
import software.aws.toolkits.resources.message
import software.aws.toolkits.telemetry.AwsTelemetry

class CredentialsFileHelpNotificationProvider : EditorNotifications.Provider<CredentialFileNotificationPanel>(), DumbAware {
    override fun getKey(): Key<CredentialFileNotificationPanel> = KEY

    override fun createNotificationPanel(file: VirtualFile, fileEditor: FileEditor, project: Project): CredentialFileNotificationPanel? {
        // Check if editor is for the config/credential file
        if (!isCredentialsFile(file)) return null

        return CredentialFileNotificationPanel(project)
    }

    private fun isCredentialsFile(file: VirtualFile): Boolean = try {
        val filePath = file.toNioPath().toAbsolutePath()
        ProfileFileLocation.configurationFilePath().toAbsolutePath() == filePath || ProfileFileLocation.credentialsFilePath().toAbsolutePath() == filePath
    } catch (e: Exception) {
        false
    }

    class CredentialFileNotificationPanel(project: Project) : EditorNotificationPanel() {
        init {
            createActionLabel(message("general.save")) {
                FileDocumentManager.getInstance().saveAllDocuments()
                AwsTelemetry.saveCredentials(project = project)
            }

            createActionLabel(message("general.help")) {
                HelpManager.getInstance().invokeHelp(HelpIds.SETUP_CREDENTIALS.id)
                AwsTelemetry.help(project = project, name = HelpIds.SETUP_CREDENTIALS.id)
            }

            text(message("credentials.file.notification"))
        }
    }

    companion object {
        /**
         * Key used to store the notification panel in an editor
         */
        val KEY = Key.create<CredentialFileNotificationPanel>("software.aws.toolkits.jetbrains.core.credentials.editor.notification.provider")
    }
}


// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import software.aws.toolkits.core.credentials.CredentialSourceId
import software.aws.toolkits.telemetry.CredentialSourceId as TelemetryCredentialSourceId

fun CredentialSourceId?.toTelemetryCredentialSourceId(): TelemetryCredentialSourceId = when (this) {
    CredentialSourceId.SharedCredentials -> TelemetryCredentialSourceId.SharedCredentials
    CredentialSourceId.SdkStore -> TelemetryCredentialSourceId.SdkStore
    CredentialSourceId.Ec2 -> TelemetryCredentialSourceId.Ec2
    CredentialSourceId.Ecs -> TelemetryCredentialSourceId.Ecs
    CredentialSourceId.EnvVars -> TelemetryCredentialSourceId.EnvVars
    else -> TelemetryCredentialSourceId.Other
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.notification.NotificationAction
import com.intellij.openapi.components.service
import com.intellij.openapi.project.Project
import software.aws.toolkits.core.credentials.CredentialIdentifier
import software.aws.toolkits.core.region.AwsRegion
import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider
import software.aws.toolkits.jetbrains.settings.AwsSettings
import software.aws.toolkits.jetbrains.settings.UseAwsCredentialRegion
import software.aws.toolkits.jetbrains.utils.notifyInfo
import software.aws.toolkits.resources.message

/**
 * Encapsulates logic for handling of regions when a new credential identifier is selected
 */
interface CredentialsRegionHandler {
    fun determineSelectedRegion(identifier: CredentialIdentifier, selectedRegion: AwsRegion?): AwsRegion?

    companion object {
        fun getInstance(project: Project): CredentialsRegionHandler = project.service()
    }
}

internal class DefaultCredentialsRegionHandler(private val project: Project) : CredentialsRegionHandler {
    override fun determineSelectedRegion(identifier: CredentialIdentifier, selectedRegion: AwsRegion?): AwsRegion? {
        val settings = AwsSettings.getInstance()
        if (settings.useDefaultCredentialRegion == UseAwsCredentialRegion.Never) {
            return selectedRegion
        }

        val regionProvider = AwsRegionProvider.getInstance()
        val defaultCredentialRegion = identifier.defaultRegionId?.let { regionProvider[it] } ?: return selectedRegion
        when {
            selectedRegion == defaultCredentialRegion -> return defaultCredentialRegion
            selectedRegion?.partitionId != defaultCredentialRegion.partitionId -> return defaultCredentialRegion
            settings.useDefaultCredentialRegion == UseAwsCredentialRegion.Always -> return defaultCredentialRegion
            settings.useDefaultCredentialRegion == UseAwsCredentialRegion.Prompt -> promptForRegionChange(defaultCredentialRegion)
        }
        return selectedRegion
    }

    private fun promptForRegionChange(defaultCredentialRegion: AwsRegion) {
        notifyInfo(
            message("aws.notification.title"),
            message("settings.credentials.prompt_for_default_region_switch", defaultCredentialRegion.id),
            project = project,
            notificationActions = listOf(
                NotificationAction.createSimpleExpiring(message("settings.credentials.prompt_for_default_region_switch.yes")) {
                    AwsConnectionManager.getInstance(project).changeRegion(defaultCredentialRegion)
                },
                NotificationAction.createSimpleExpiring(message("settings.credentials.prompt_for_default_region_switch.always")) {
                    AwsSettings.getInstance().useDefaultCredentialRegion = UseAwsCredentialRegion.Always
                    AwsConnectionManager.getInstance(project).changeRegion(defaultCredentialRegion)
                },
                NotificationAction.createSimpleExpiring(message("settings.credentials.prompt_for_default_region_switch.never")) {
                    AwsSettings.getInstance().useDefaultCredentialRegion = UseAwsCredentialRegion.Never
                }
            )
        )
    }
}


// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.openapi.actionSystem.ActionManager
import com.intellij.openapi.project.Project
import software.aws.toolkits.jetbrains.utils.createNotificationExpiringAction
import software.aws.toolkits.jetbrains.utils.createShowMoreInfoDialogAction
import software.aws.toolkits.jetbrains.utils.notifyWarn
import software.aws.toolkits.resources.message

class CredentialStatusNotification(private val project: Project) : ConnectionSettingsStateChangeNotifier {
    private val actionManager = ActionManager.getInstance()
    override fun settingsStateChanged(newState: ConnectionState) {
        if (newState is ConnectionState.InvalidConnection) {
            val title = message("credentials.invalid.title")
            val message = message("credentials.invalid.description")

            notifyWarn(
                project = project,
                title = title,
                content = message,
                notificationActions = listOf(
                    createShowMoreInfoDialogAction(
                        message("credentials.invalid.more_info"),
                        title,
                        message,
                        newState.displayMessage
                    ),
                    createNotificationExpiringAction(actionManager.getAction("aws.settings.upsertCredentials")),
                    createNotificationExpiringAction(RefreshConnectionAction(message("settings.retry")))
                )
            )
        }
    }
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import software.aws.toolkits.core.credentials.CredentialType
import software.aws.toolkits.telemetry.CredentialType as TelemetryCredentialType

fun CredentialType?.toTelemetryType(): TelemetryCredentialType = when (this) {
    CredentialType.StaticProfile -> TelemetryCredentialType.StaticProfile
    CredentialType.StaticSessionProfile -> TelemetryCredentialType.StaticSessionProfile
    CredentialType.CredentialProcessProfile -> TelemetryCredentialType.CredentialProcessProfile
    CredentialType.AssumeRoleProfile -> TelemetryCredentialType.AssumeRoleProfile
    CredentialType.AssumeMfaRoleProfile -> TelemetryCredentialType.AssumeMfaRoleProfile
    CredentialType.SsoProfile -> TelemetryCredentialType.SsoProfile
    CredentialType.Ec2Metadata -> TelemetryCredentialType.Ec2Metadata
    CredentialType.EcsMetadata -> TelemetryCredentialType.EcsMetatdata
    null -> TelemetryCredentialType.Other
}


// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.icons.AllIcons
import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.actionSystem.PlatformDataKeys
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.fileEditor.FileEditorManager
import com.intellij.openapi.fileEditor.OpenFileDescriptor
import com.intellij.openapi.fileTypes.FileTypes
import com.intellij.openapi.fileTypes.ex.FileTypeManagerEx
import com.intellij.openapi.project.DumbAware
import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.Messages
import com.intellij.openapi.vfs.LocalFileSystem
import org.jetbrains.annotations.TestOnly
import software.amazon.awssdk.profiles.ProfileFileLocation
import software.aws.toolkits.core.utils.createParentDirectories
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.touch
import software.aws.toolkits.core.utils.tryDirOp
import software.aws.toolkits.core.utils.tryFileOp
import software.aws.toolkits.core.utils.writeText
import software.aws.toolkits.resources.message
import software.aws.toolkits.telemetry.AwsTelemetry
import java.io.File

class CreateOrUpdateCredentialProfilesAction @TestOnly constructor(
    private val writer: ConfigFileWriter,
    private val configFile: File,
    private val credentialsFile: File
) : AnAction(message("configure.toolkit.upsert_credentials.action")), DumbAware {
    @Suppress("unused")
    constructor() : this(
        DefaultConfigFileWriter,
        ProfileFileLocation.configurationFilePath().toFile(),
        ProfileFileLocation.credentialsFilePath().toFile()
    )

    private val localFileSystem = LocalFileSystem.getInstance()

    override fun actionPerformed(e: AnActionEvent) {
        val project = e.getRequiredData(PlatformDataKeys.PROJECT)

        // if both config and credential files do not exist, create a new config file
        if (!configFile.exists() && !credentialsFile.exists()) {
            if (confirm(project, configFile)) {
                try {
                    writer.createFile(configFile)
                } finally {
                    AwsTelemetry.createCredentials(project)
                }
            } else {
                return
            }
        }

        // open both config and credential files, if they exist
        // credential file is opened last since it takes precedence over the config file
        val virtualFiles = listOf(configFile, credentialsFile).filter { it.exists() }.map {
            localFileSystem.refreshAndFindFileByIoFile(it) ?: throw RuntimeException(message("credentials.could_not_open", it))
        }

        val fileEditorManager = FileEditorManager.getInstance(project)

        localFileSystem.refreshFiles(virtualFiles, false, false) {
            virtualFiles.forEach {
                if (it.fileType == FileTypes.UNKNOWN) {
                    ApplicationManager.getApplication().runWriteAction {
                        FileTypeManagerEx.getInstanceEx().associatePattern(
                            FileTypes.PLAIN_TEXT,
                            it.name
                        )
                    }
                }

                if (fileEditorManager.openTextEditor(OpenFileDescriptor(project, it), true) == null) {
                    AwsTelemetry.openCredentials(project, success = false)
                    throw RuntimeException(message("credentials.could_not_open", it))
                }
                AwsTelemetry.openCredentials(project, success = true)
            }
        }
    }

    private fun confirm(project: Project, file: File): Boolean = Messages.showOkCancelDialog(
        project,
        message("configure.toolkit.upsert_credentials.confirm_file_create", file),
        message("configure.toolkit.upsert_credentials.confirm_file_create.title"),
        message("configure.toolkit.upsert_credentials.confirm_file_create.okay"),
        Messages.getCancelButton(),
        AllIcons.General.QuestionDialog,
        null
    ) == Messages.OK
}

interface ConfigFileWriter {
    fun createFile(file: File)
}

object DefaultConfigFileWriter : ConfigFileWriter {
    val TEMPLATE =
        """
        # Amazon Web Services Config File used by AWS CLI, SDKs, and tools
        # This file was created by the AWS Toolkit for JetBrains plugin.
        #
        # Your AWS credentials are represented by access keys associated with IAM users.
        # For information about how to create and manage AWS access keys for a user, see:
        # https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_access-keys.html
        #
        # This config file can store multiple access keys by placing each one in a
        # named "profile". For information about how to change the access keys in a
        # profile or to add a new profile with a different access key, see:
        # https://docs.aws.amazon.com/cli/latest/userguide/cli-config-files.html
        #
        # If both a credential and config file exists, the values in the credential file
        # take precedence

        [default]
        # The access key and secret key pair identify your account and grant access to AWS.
        aws_access_key_id = [accessKey]
        # Treat your secret key like a password. Never share your secret key with anyone. Do
        # not post it in online forums, or store it in a source control system. If your secret
        # key is ever disclosed, immediately use IAM to delete the access key and secret key
        # and create a new key pair. Then, update this file with the replacement key details.
        aws_secret_access_key = [secretKey]

        # [profile user1]
        # aws_access_key_id = [accessKey1]
        # aws_secret_access_key = [secretKey1]
        """.trimIndent()

    override fun createFile(file: File) {
        val path = file.toPath()
        path.tryDirOp(LOG) { createParentDirectories() }

        path.tryFileOp(LOG) {
            touch(restrictToOwner = true)
            writeText(TEMPLATE)
        }
    }

    private val LOG = getLogger<DefaultConfigFileWriter>()
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.openapi.components.PersistentStateComponent
import com.intellij.openapi.components.State
import com.intellij.openapi.components.Storage
import com.intellij.openapi.project.Project
import kotlinx.coroutines.launch
import software.aws.toolkits.core.utils.tryOrNull
import software.aws.toolkits.jetbrains.core.coroutines.disposableCoroutineScope
import software.aws.toolkits.jetbrains.core.credentials.profiles.DEFAULT_PROFILE_ID
import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider

data class ConnectionSettingsState(
    var activeProfile: String? = null,
    var activeRegion: String? = null,
    var recentlyUsedProfiles: List<String> = mutableListOf(),
    var recentlyUsedRegions: List<String> = mutableListOf()
)

@State(name = "accountSettings", storages = [Storage("aws.xml")])
class DefaultAwsConnectionManager(project: Project) :
    AwsConnectionManager(project),
    PersistentStateComponent<ConnectionSettingsState> {
    private val coroutineScope = disposableCoroutineScope(this)

    override fun getState(): ConnectionSettingsState = ConnectionSettingsState(
        activeProfile = selectedCredentialIdentifier?.id,
        activeRegion = selectedRegion?.id,
        recentlyUsedProfiles = recentlyUsedProfiles.elements(),
        recentlyUsedRegions = recentlyUsedRegions.elements()
    )

    override fun noStateLoaded() {
        loadState(ConnectionSettingsState())
    }

    override fun loadState(state: ConnectionSettingsState) {
        // This can be called more than once, so we need to re-do our init sequence
        connectionState = ConnectionState.InitializingToolkit

        // Load reversed so that oldest is as the bottom
        state.recentlyUsedRegions.reversed()
            .forEach { recentlyUsedRegions.add(it) }

        state.recentlyUsedProfiles.reversed()
            .forEach { recentlyUsedProfiles.add(it) }

        // Load all the initial state on BG thread, so we don't block the UI or loading of other components
        coroutineScope.launch {
            val credentialId = state.activeProfile ?: DEFAULT_PROFILE_ID
            val credentials = tryOrNull {
                CredentialManager.getInstance().getCredentialIdentifierById(credentialId)
            }

            val regionId = state.activeRegion ?: credentials?.defaultRegionId ?: AwsRegionProvider.getInstance().defaultRegion().id
            val region = AwsRegionProvider.getInstance().allRegions()[regionId]

            changeConnectionSettings(credentials, region)
        }
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.openapi.Disposable
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.components.PersistentStateComponent
import com.intellij.openapi.components.State
import com.intellij.openapi.components.Storage
import com.intellij.openapi.util.Disposer
import software.aws.toolkits.core.utils.error
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.info
import software.aws.toolkits.core.utils.warn
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener

// TODO: unify with CredentialManager
@State(name = "authManager", storages = [Storage("aws.xml")])
class DefaultToolkitAuthManager : ToolkitAuthManager, PersistentStateComponent<ToolkitAuthManagerState>, Disposable {
    private var state = ToolkitAuthManagerState()
    private val connections = linkedSetOf<ToolkitConnection>()
    private val transientConnections = let {
        val factoryConnections = mutableListOf<ToolkitConnection>()
        ToolkitStartupAuthFactory.EP_NAME.forEachExtensionSafe { factory ->
            factoryConnections.addAll(
                factory.buildConnections().also { connections ->
                    LOG.info { "Found transient connections from $factory: ${connections.map { it.toString() }}" }
                }
            )
        }

        factoryConnections.toList()
    }

    override fun listConnections(): List<ToolkitConnection> = connections.toList() + transientConnections

    override fun createConnection(profile: AuthProfile): ToolkitConnection {
        val connection = connectionFromProfile(profile)
        connections.firstOrNull { it.id == connection.id }?.let {
            LOG.warn { "$it already exists in connection list" }
            if (connection is Disposable) {
                Disposer.dispose(connection)
            }

            return it
        }

        connections.add(connection)
        return connection
    }

    private fun deleteConnection(predicate: (ToolkitConnection) -> Boolean) {
        connections.removeAll { connection ->
            predicate(connection).also {
                if (it && connection is Disposable) {
                    ApplicationManager.getApplication().messageBus.syncPublisher(BearerTokenProviderListener.TOPIC)
                        .invalidate(connection.id)
                    Disposer.dispose(connection)
                }
            }
        }
    }

    override fun deleteConnection(connection: ToolkitConnection) {
        deleteConnection { it == connection }
    }

    override fun deleteConnection(connectionId: String) {
        deleteConnection { it.id == connectionId }
    }

    override fun getConnection(connectionId: String) = listConnections().firstOrNull { it.id == connectionId }

    override fun getState(): ToolkitAuthManagerState? {
        val data = connections.mapNotNull {
            when (it) {
                is ManagedBearerSsoConnection -> {
                    ManagedSsoProfile(
                        startUrl = it.startUrl,
                        ssoRegion = it.region,
                        scopes = it.scopes
                    )
                }

                else -> {
                    LOG.error { "Couldn't serialize $it" }
                    null
                }
            }
        }

        state.ssoProfiles = data

        return state
    }

    override fun loadState(state: ToolkitAuthManagerState) {
        this.state = state
        val newConnections = linkedSetOf(*state.ssoProfiles.toTypedArray()).filterNotNull().map {
            connectionFromProfile(it)
        }

        if (newConnections.size != state.ssoProfiles.size) {
            LOG.warn { "Persisted state had duplicate profiles" }
        }

        connections.clear()
        connections.addAll(newConnections)
    }

    override fun dispose() {
        listConnections().forEach {
            if (it is Disposable) {
                Disposer.dispose(it)
            }
        }
    }

    private fun connectionFromProfile(profile: AuthProfile): ToolkitConnection = when (profile) {
        is ManagedSsoProfile -> {
            ManagedBearerSsoConnection(
                startUrl = profile.startUrl,
                region = profile.ssoRegion,
                scopes = profile.scopes
            )
        }

        is DiskSsoSessionProfile -> DiskSsoSessionConnection(sessionProfileName = profile.profileName, region = profile.ssoRegion)
    }

    companion object {
        private val LOG = getLogger<DefaultToolkitAuthManager>()
    }
}

data class ToolkitAuthManagerState(
    // TODO: can't figure out how to make deserializer work with polymorphic types
    var ssoProfiles: List<ManagedSsoProfile> = emptyList()
)


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.ide.ActivityTracker
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.components.PersistentStateComponent
import com.intellij.openapi.components.State
import com.intellij.openapi.components.Storage
import com.intellij.openapi.project.Project
import software.aws.toolkits.jetbrains.core.credentials.pinning.ConnectionPinningManager
import software.aws.toolkits.jetbrains.core.credentials.pinning.FeatureWithPinnedConnection
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener

// TODO: unify with AwsConnectionManager
@State(name = "connectionManager", storages = [Storage("aws.xml")])
class DefaultToolkitConnectionManager : ToolkitConnectionManager, PersistentStateComponent<ToolkitConnectionManagerState> {
    init {
        ApplicationManager.getApplication().messageBus.connect(this).subscribe(
            BearerTokenProviderListener.TOPIC,
            object : BearerTokenProviderListener {
                override fun invalidate(providerId: String) {
                    if (activeConnection()?.id == providerId) {
                        switchConnection(null)
                        ActivityTracker.getInstance().inc()
                    }
                }
            }
        )
    }
    private val project: Project?
    constructor(project: Project) {
        this.project = project
    }
    constructor() {
        this.project = null
    }

    private var connection: ToolkitConnection? = null

    private val pinningManager: ConnectionPinningManager = ConnectionPinningManager.getInstance()

    private val defaultConnection: ToolkitConnection?
        get() {
            if (CredentialManager.getInstance().getCredentialIdentifiers().isNotEmpty() && project != null) {
                return AwsConnectionManagerConnection(project)
            }

            ToolkitAuthManager.getInstance().listConnections().firstOrNull()?.let {
                return it
            }

            return null
        }

    @Synchronized
    override fun activeConnection() = connection ?: defaultConnection

    @Synchronized
    override fun activeConnectionForFeature(feature: FeatureWithPinnedConnection): ToolkitConnection? {
        val pinnedConnection = pinningManager.getPinnedConnection(feature)
        if (pinnedConnection != null) {
            return pinnedConnection
        }

        return connection?.let {
            if (feature.supportsConnectionType(it)) {
                return it
            }

            null
        } ?: defaultConnection?.let {
            if (feature.supportsConnectionType(it)) {
                return it
            }

            null
        }
    }

    override fun getState() = ToolkitConnectionManagerState(
        connection?.id
    )

    override fun loadState(state: ToolkitConnectionManagerState) {
        state.activeConnectionId?.let {
            val idSegments = it.split(";")
            val activeConnectionIdWithRegion =
                if (idSegments.size == 2) {
                    "${idSegments[0]};us-east-1;${idSegments[1]}"
                } else {
                    it
                }
            connection = ToolkitAuthManager.getInstance().getConnection(activeConnectionIdWithRegion)
        }
    }

    @Synchronized
    override fun switchConnection(newConnection: ToolkitConnection?) {
        val oldConnection = this.connection

        if (oldConnection != newConnection) {
            val application = ApplicationManager.getApplication()
            this.connection = newConnection

            if (newConnection != null) {
                val featuresToPin = mutableListOf<FeatureWithPinnedConnection>()
                FeatureWithPinnedConnection.EP_NAME.forEachExtensionSafe {
                    if (!pinningManager.isFeaturePinned(it) &&
                        (
                            (oldConnection == null && it.supportsConnectionType(newConnection)) ||
                                (oldConnection != null && it.supportsConnectionType(oldConnection) != it.supportsConnectionType(newConnection))
                            )
                    ) {
                        featuresToPin.add(it)
                    }
                }

                if (featuresToPin.isNotEmpty()) {
                    application.executeOnPooledThread {
                        pinningManager.maybePinFeatures(oldConnection, newConnection, featuresToPin)
                        application.messageBus.syncPublisher(ToolkitConnectionManagerListener.TOPIC).activeConnectionChanged(newConnection)
                    }
                }
            }

            application.messageBus.syncPublisher(ToolkitConnectionManagerListener.TOPIC).activeConnectionChanged(newConnection)
        }
    }
}

data class ToolkitConnectionManagerState(
    var activeConnectionId: String? = null
)


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.openapi.actionSystem.AnAction
import software.aws.toolkits.core.credentials.CredentialIdentifier

/**
 * Interface that indicates that [CredentialIdentifier] may require interaction from a user before they can be used
 */
interface InteractiveCredential : CredentialIdentifier {
    val userActionDisplayMessage: String
    val userActionShortDisplayMessage: String get() = userActionDisplayMessage

    val userAction: AnAction

    /**
     * Determines if user action is required at this time (e.g. may check expiry of cookies, etc)
     */
    fun userActionRequired(): Boolean
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.progress.ProcessCanceledException
import com.intellij.openapi.ui.Messages
import software.aws.toolkits.jetbrains.utils.computeOnEdt
import software.aws.toolkits.resources.message

interface MfaRequiredInteractiveCredentials : InteractiveCredential {
    override val userActionDisplayMessage: String get() = message("credentials.mfa.display", displayName)
    override val userActionShortDisplayMessage: String get() = message("credentials.mfa.display.short")

    override val userAction: AnAction get() = RefreshConnectionAction(message("credentials.mfa.action"))

    override fun userActionRequired(): Boolean = true
}

fun promptForMfaToken(name: String, mfaSerial: String): String = computeOnEdt {
    Messages.showInputDialog(
        message("credentials.mfa.message", mfaSerial),
        message("credentials.mfa.title", name),
        null
    ) ?: throw ProcessCanceledException(IllegalStateException("MFA challenge is required"))
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.icons.AllIcons
import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.project.DumbAware
import kotlinx.coroutines.launch
import software.aws.toolkits.jetbrains.core.AwsResourceCache
import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
import software.aws.toolkits.resources.message
import software.aws.toolkits.telemetry.AwsTelemetry

class RefreshConnectionAction(text: String = message("settings.refresh.description")) :
    AnAction(text, null, AllIcons.Actions.Refresh),
    DumbAware {
    override fun update(e: AnActionEvent) {
        val project = e.project ?: return
        e.presentation.isEnabled = when (val state = AwsConnectionManager.getInstance(project).connectionState) {
            is ConnectionState.IncompleteConfiguration -> false
            else -> state.isTerminal
        }
    }

    override fun actionPerformed(e: AnActionEvent) {
        val project = e.project ?: return
        AwsTelemetry.refreshExplorer(project)
        val scope = projectCoroutineScope(project)
        scope.launch { AwsResourceCache.getInstance().clear() }
        AwsConnectionManager.getInstance(project).refreshConnectionState()
    }
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.openapi.actionSystem.AnAction
import software.aws.toolkits.jetbrains.core.credentials.sso.DiskCache
import software.aws.toolkits.jetbrains.core.credentials.sso.SsoCache
import software.aws.toolkits.resources.message

/**
 * Shared disk cache for SSO for the IDE
 */
val diskCache by lazy { DiskCache() }

interface SsoRequiredInteractiveCredentials : InteractiveCredential {
    val ssoCache: SsoCache
    val ssoUrl: String

    override val userActionDisplayMessage: String get() = message("credentials.sso.display", displayName)
    override val userActionShortDisplayMessage: String get() = message("credentials.sso.display.short")

    override val userAction: AnAction get() = RefreshConnectionAction(message("credentials.sso.action"))

    override fun userActionRequired(): Boolean = ssoCache.loadAccessToken(ssoUrl) == null
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.openapi.Disposable
import com.intellij.openapi.actionSystem.ActionManager
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.actionSystem.DataContext
import com.intellij.openapi.actionSystem.PlatformDataKeys
import com.intellij.openapi.application.ModalityState
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.progress.ProcessCanceledException
import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.DialogPanel
import com.intellij.openapi.ui.DialogWrapper
import com.intellij.ui.CollectionComboBoxModel
import com.intellij.ui.components.JBRadioButton
import com.intellij.ui.dsl.builder.COLUMNS_MEDIUM
import com.intellij.ui.dsl.builder.Cell
import com.intellij.ui.dsl.builder.RowLayout
import com.intellij.ui.dsl.builder.TopGap
import com.intellij.ui.dsl.builder.bind
import com.intellij.ui.dsl.builder.bindItem
import com.intellij.ui.dsl.builder.bindText
import com.intellij.ui.dsl.builder.columns
import com.intellij.ui.dsl.builder.panel
import com.intellij.ui.dsl.builder.selected
import com.intellij.ui.dsl.builder.toNullableProperty
import com.intellij.util.containers.nullize
import software.amazon.awssdk.services.ssooidc.model.InvalidGrantException
import software.amazon.awssdk.services.ssooidc.model.InvalidRequestException
import software.amazon.awssdk.services.ssooidc.model.SsoOidcException
import software.aws.toolkits.core.credentials.DEFAULT_SSO_REGION
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.info
import software.aws.toolkits.core.utils.warn
import software.aws.toolkits.jetbrains.ToolkitPlaces
import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_URL
import software.aws.toolkits.jetbrains.core.help.HelpIds
import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider
import software.aws.toolkits.jetbrains.utils.runUnderProgressIfNeeded
import software.aws.toolkits.resources.message
import software.aws.toolkits.telemetry.UiTelemetry
import java.io.IOException
import javax.swing.JComponent

data class ConnectionDialogCustomizer(
    val title: String? = null,
    val header: String? = null,
    val helpId: HelpIds? = null,
    val replaceIamComment: String? = null,
    val scopes: List<String>? = null,
    val startUrl: String? = null,
    val region: String? = null,
    val errorMsg: String? = null
)

open class ToolkitAddConnectionDialog(
    private val project: Project,
    private val customizer: ConnectionDialogCustomizer? = null,
) : DialogWrapper(project), Disposable {
    private data class Modal(
        var loginType: LoginOptions,
        var startUrl: String,
        var region: String
    ) {
        constructor() : this(LoginOptions.AWS_BUILDER_ID, SONO_URL, DEFAULT_SSO_REGION)
    }

    private enum class LoginOptions {
        AWS_BUILDER_ID,
        SSO,
        IAM
    }

    private var modal: Modal = Modal()
    private val panel: DialogPanel by lazy { createPanel() }
    private val regions = AwsRegionProvider.getInstance().allRegions().values.filter { it.partitionId == "aws" }.map { it.id }

    constructor(project: Project, connection: ToolkitConnection? = null, customizer: ConnectionDialogCustomizer?) : this(project, customizer) {
        // Fill in login metadata for users to reauthenticate
        connection?.let {
            if (it is ManagedBearerSsoConnection) {
                when (it.startUrl) {
                    SONO_URL -> run { modal.loginType = LoginOptions.AWS_BUILDER_ID }
                    else -> run {
                        modal.loginType = LoginOptions.SSO
                        modal.region = it.region
                        modal.startUrl = it.startUrl
                    }
                }
            }
        }
    }

    init {
        title = customizer?.title ?: message("toolkit.login.dialog.title")
        setOKButtonText(message("toolkit.login.dialog.connect_button"))

        modal.apply {
            this.startUrl = customizer?.startUrl.orEmpty()
            this.region = customizer?.region ?: DEFAULT_SSO_REGION
            this.loginType = if (startUrl.isEmpty()) LoginOptions.AWS_BUILDER_ID else LoginOptions.SSO
        }

        init()
        // this must happen after init()
        customizer?.errorMsg?.let { setErrorText(it) }
    }

    override fun createCenterPanel(): JComponent = panel

    override fun getHelpId(): String = (customizer?.helpId ?: HelpIds.TOOLKIT_ADD_CONNECTIONS_DIALOG).id

    override fun doHelpAction() {
        UiTelemetry.click(project, "connection_help")
        super.doHelpAction()
    }

    override fun doCancelAction() {
        UiTelemetry.click(project, "connection_optionescapecancel")
        super.doCancelAction()
    }

    override fun doOKAction() {
        if (!isOKActionEnabled) {
            return
        }
        this.panel.apply()
        setOKButtonText(message("toolkit.login.dialog.connect_inprogress"))

        val loginType = modal.loginType
        val startUrl = if (loginType == LoginOptions.AWS_BUILDER_ID) SONO_URL else modal.startUrl
        val region = modal.region
        val progressIndicatorTitle = when (loginType) {
            LoginOptions.AWS_BUILDER_ID -> "AWS Builder ID"
            LoginOptions.SSO -> "IAM Identity Center"
            else -> error("User should not be able to choose option other than above ")
        }

        try {
            close(OK_EXIT_CODE)
            runUnderProgressIfNeeded(project, "Login: $progressIndicatorTitle", true) {
                // Edge case when user choose SSO but enter AWS Builder ID url
                if (loginType == LoginOptions.SSO && startUrl == SONO_URL) {
                    error("User should not perform Identity Center login with AWS Builder ID url")
                }

                val scopes = customizer?.scopes?.nullize() ?: listOf("sso:account:access")

                LOG.info { "Try to fetch credential with: $modal" }

                loginSso(project, startUrl, region, scopes)
            }
        } catch (e: Exception) {
            val message = when (e) {
                is IllegalStateException -> e.message ?: message("general.unknown_error")
                is ProcessCanceledException -> message("codewhisperer.credential.login.dialog.exception.cancel_login")
                is InvalidGrantException -> message("codewhisperer.credential.login.exception.invalid_grant")
                is InvalidRequestException -> message("codewhisperer.credential.login.exception.invalid_input")
                is SsoOidcException -> message("codewhisperer.credential.login.exception.general.oidc")
                else -> {
                    val baseMessage = when (e) {
                        is IOException -> "codewhisperer.credential.login.exception.io"
                        else -> "codewhisperer.credential.login.exception.general"
                    }

                    message(baseMessage, "${e.javaClass.name}: ${e.message}")
                }
            }

            LOG.warn(e) { "Failed to fetch credential with: $modal; reason: $message" }

            runInEdt(ModalityState.any()) {
                ToolkitAddConnectionDialog(
                    project,
                    (customizer ?: ConnectionDialogCustomizer()).copy(
                        startUrl = if (startUrl == SONO_URL) "" else startUrl, // TODO: fix this, ugly
                        region = region,
                        errorMsg = message
                    )
                ).showAndGet()
            }
        } finally {
            val credType = when (loginType) {
                LoginOptions.AWS_BUILDER_ID -> "connection_optionBuilderID"
                LoginOptions.SSO -> "connection_optionSSO"
                else -> null
            }

            credType?.let { UiTelemetry.click(project, it) }
        }
    }

    override fun dispose() {
        super.dispose()
    }

    private fun createPanel() = panel {
        row {
            label(customizer?.header ?: message("toolkit.login.dialog.label"))
                .bold()
        }

        buttonsGroup {
            // AWS Builder ID
            row {
                radioButton(message("toolkit.login.dialog.aws_builder_id.title"), LoginOptions.AWS_BUILDER_ID)
                    .comment(
                        message("toolkit.login.dialog.aws_builder_id.comment"),
                        commentMaxLength
                    )
            }

            // SSO
            lateinit var ssoRadioButton: Cell<JBRadioButton>
            row {
                ssoRadioButton = radioButton(message("toolkit.login.dialog.sso.title"), LoginOptions.SSO)
                    .comment(
                        message("toolkit.login.dialog.sso.comment"),
                        commentMaxLength
                    )
            }.topGap(TopGap.MEDIUM)

            indent {
                row {
                    panel {
                        row {
                            label(message("toolkit.login.dialog.sso.text_field.start_url"))
                        }
                        row {
                            textField().apply {
                                columns(COLUMNS_MEDIUM)
                                bindText(modal::startUrl)
                            }
                        }
                            .enabledIf(ssoRadioButton.selected)
                            .layout(RowLayout.INDEPENDENT)
                    }

                    panel {
                        row(message("toolkit.login.dialog.sso.text_field.region")) {}
                        row {
                            comboBox(CollectionComboBoxModel(regions, DEFAULT_SSO_REGION))
                                .bindItem(modal::region.toNullableProperty())
                        }
                            .enabledIf(ssoRadioButton.selected)
                            .layout(RowLayout.INDEPENDENT)
                    }
                }
            }
            // IAM
            row {
                radioButton(message("toolkit.login.dialog.iam.title"), LoginOptions.IAM)
                    .enabled(false)
                    .comment(
                        customizer?.replaceIamComment ?: "<a>${message("configure.toolkit.upsert_credentials.action")}</a>",
                        commentMaxLength
                    ) { hyperlinkEvent ->
                        close(OK_EXIT_CODE)
                        val actionEvent = AnActionEvent.createFromInputEvent(
                            hyperlinkEvent.inputEvent,
                            ToolkitPlaces.ADD_CONNECTION_DIALOG,
                            null,
                            DataContext { if (PlatformDataKeys.PROJECT.`is`(it)) project else null }
                        )
                        ActionManager.getInstance().getAction("aws.settings.upsertCredentials").actionPerformed(actionEvent)
                    }
            }.topGap(TopGap.MEDIUM)
        }.bind(modal::loginType)

        separator().topGap(TopGap.MEDIUM)
    }

    private companion object {
        private val LOG = getLogger<ToolkitAddConnectionDialog>()
        const val commentMaxLength = 85
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.openapi.Disposable
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.components.service
import com.intellij.openapi.extensions.ExtensionPointName
import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.MessageDialogBuilder
import org.jetbrains.annotations.VisibleForTesting
import software.amazon.awssdk.services.ssooidc.model.SsoOidcException
import software.aws.toolkits.core.ClientConnectionSettings
import software.aws.toolkits.core.ConnectionSettings
import software.aws.toolkits.core.TokenConnectionSettings
import software.aws.toolkits.core.credentials.ToolkitBearerTokenProvider
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.info
import software.aws.toolkits.core.utils.warn
import software.aws.toolkits.jetbrains.core.credentials.pinning.FeatureWithPinnedConnection
import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenAuthState
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
import software.aws.toolkits.jetbrains.utils.computeOnEdt
import software.aws.toolkits.jetbrains.utils.runUnderProgressIfNeeded
import software.aws.toolkits.resources.message

sealed interface ToolkitConnection {
    val id: String
    val label: String

    fun getConnectionSettings(): ClientConnectionSettings<*>
}

interface AwsCredentialConnection : ToolkitConnection {
    override fun getConnectionSettings(): ConnectionSettings
}

interface AwsBearerTokenConnection : ToolkitConnection {
    override fun getConnectionSettings(): TokenConnectionSettings
}

interface BearerSsoConnection : AwsBearerTokenConnection {
    val scopes: List<String>
}

sealed interface AuthProfile

data class ManagedSsoProfile(
    var ssoRegion: String = "",
    var startUrl: String = "",
    var scopes: List<String> = emptyList()
) : AuthProfile

data class DiskSsoSessionProfile(
    var profileName: String = "",
    var ssoRegion: String = ""
) : AuthProfile

/**
 * Used to contribute connections to [ToolkitAuthManager] on service initialization
 */
interface ToolkitStartupAuthFactory {
    fun buildConnections(): List<ToolkitConnection>

    companion object {
        val EP_NAME = ExtensionPointName.create<ToolkitStartupAuthFactory>("aws.toolkit.startupAuthFactory")
    }
}

interface ToolkitAuthManager {
    fun listConnections(): List<ToolkitConnection>

    fun createConnection(profile: AuthProfile): ToolkitConnection

    fun deleteConnection(connection: ToolkitConnection)
    fun deleteConnection(connectionId: String)

    fun getConnection(connectionId: String): ToolkitConnection?

    companion object {
        fun getInstance() = service<ToolkitAuthManager>()
    }
}

interface ToolkitConnectionManager : Disposable {
    fun activeConnection(): ToolkitConnection?

    fun activeConnectionForFeature(feature: FeatureWithPinnedConnection): ToolkitConnection?

    fun switchConnection(newConnection: ToolkitConnection?)

    override fun dispose() {}

    companion object {
        fun getInstance(project: Project?) = project?.let { it.service<ToolkitConnectionManager>() } ?: service()
    }
}

/**
 * Individual service should subscribe [ToolkitConnectionManagerListener.TOPIC] to fire their service activation / UX update
 */

fun loginSso(project: Project?, startUrl: String, region: String, requestedScopes: List<String>): BearerTokenProvider {
    val connectionId = ToolkitBearerTokenProvider.ssoIdentifier(startUrl, region)

    val manager = ToolkitAuthManager.getInstance()
    val allScopes = requestedScopes.toMutableList()
    return manager.getConnection(connectionId)?.let { connection ->
        val logger = getLogger<ToolkitAuthManager>()
        // requested Builder ID, but one already exists
        // TBD: do we do this for regular SSO too?
        if (connection.isSono() && connection is BearerSsoConnection && requestedScopes.all { it in connection.scopes }) {
            val signOut = computeOnEdt {
                MessageDialogBuilder.yesNo(
                    message("toolkit.login.aws_builder_id.already_connected.title"),
                    message("toolkit.login.aws_builder_id.already_connected.message")
                )
                    .yesText(message("toolkit.login.aws_builder_id.already_connected.reconnect"))
                    .noText(message("toolkit.login.aws_builder_id.already_connected.cancel"))
                    .ask(project)
            }

            if (signOut) {
                logger.info {
                    "Forcing reauth on ${connection.id} since user requested Builder ID while already connected to Builder ID"
                }

                logoutFromSsoConnection(project, connection as AwsBearerTokenConnection)
                return@let null
            }
        }

        // There is an existing connection we can use
        if (connection is BearerSsoConnection && !requestedScopes.all { it in connection.scopes }) {
            allScopes.addAll(connection.scopes)

            logger.info {
                """
                    Forcing reauth on ${connection.id} since requested scopes ($requestedScopes)
                    are not a complete subset of current scopes (${connection.scopes})
                """.trimIndent()
            }
            logoutFromSsoConnection(project, connection as AwsBearerTokenConnection)
            // can't reuse since requested scopes are not in current connection. forcing reauth
            return@let null
        }

        // For the case when the existing connection is in invalid state, we need to re-auth
        if (connection is AwsBearerTokenConnection) {
            return reauthConnection(project, connection)
        }

        null
    } ?: run {
        // No existing connection, start from scratch
        val connection = manager.createConnection(
            ManagedSsoProfile(
                region,
                startUrl,
                allScopes.toSet().toList()
            )
        )

        try {
            reauthConnection(project, connection)
        } catch (e: Exception) {
            manager.deleteConnection(connection)
            throw e
        }
    }
}

@VisibleForTesting
internal fun reauthConnection(project: Project?, connection: ToolkitConnection): BearerTokenProvider {
    val provider = reauthProviderIfNeeded(project, connection)

    ToolkitConnectionManager.getInstance(project).switchConnection(connection)

    return provider
}

fun logoutFromSsoConnection(project: Project?, connection: AwsBearerTokenConnection, callback: () -> Unit = {}) {
    try {
        ApplicationManager.getApplication().messageBus.syncPublisher(BearerTokenProviderListener.TOPIC).invalidate(connection.id)
        ToolkitAuthManager.getInstance().deleteConnection(connection.id)
        project?.let { ToolkitConnectionManager.getInstance(it).switchConnection(null) }
    } finally {
        callback()
    }
}

fun lazyGetUnauthedBearerConnections() =
    ToolkitAuthManager.getInstance().listConnections().filterIsInstance<AwsBearerTokenConnection>().filter {
        it.lazyIsUnauthedBearerConnection()
    }

fun AwsBearerTokenConnection.lazyIsUnauthedBearerConnection(): Boolean {
    val provider = (getConnectionSettings().tokenProvider.delegate as? BearerTokenProvider)

    if (provider != null) {
        if (provider.currentToken() == null) {
            // provider is unauthed if no token
            return true
        }

        // or state is not authorized
        return provider.state() != BearerTokenAuthState.AUTHORIZED
    }

    // not a bearer token provider
    return false
}

fun reauthProviderIfNeeded(project: Project?, connection: ToolkitConnection): BearerTokenProvider {
    val tokenProvider = (connection.getConnectionSettings() as TokenConnectionSettings).tokenProvider.delegate as BearerTokenProvider

    return reauthProviderIfNeeded(project, tokenProvider)
}

fun reauthProviderIfNeeded(project: Project?, tokenProvider: BearerTokenProvider): BearerTokenProvider {
    maybeReauthProviderIfNeeded(project, tokenProvider) {
        runUnderProgressIfNeeded(project, message("credentials.sono.login.pending"), true) {
            tokenProvider.reauthenticate()
        }
    }

    return tokenProvider
}

// Return true if need to re-auth, false otherwise
fun maybeReauthProviderIfNeeded(project: Project?, tokenProvider: BearerTokenProvider, onReauthRequired: (SsoOidcException?) -> Any): Boolean {
    val state = tokenProvider.state()
    when (state) {
        BearerTokenAuthState.NOT_AUTHENTICATED -> {
            getLogger<ToolkitAuthManager>().info { "Token provider NOT_AUTHENTICATED, requesting login" }
            onReauthRequired(null)
            return true
        }

        BearerTokenAuthState.NEEDS_REFRESH -> {
            try {
                return runUnderProgressIfNeeded(project, message("credentials.sono.login.refreshing"), true) {
                    tokenProvider.resolveToken()
                    BearerTokenProviderListener.notifyCredUpdate(tokenProvider.id)
                    return@runUnderProgressIfNeeded false
                }
            } catch (e: SsoOidcException) {
                getLogger<ToolkitAuthManager>().warn(e) { "Redriving AWS Builder ID login flow since token could not be refreshed" }
                onReauthRequired(e)
                return true
            }
        }

        BearerTokenAuthState.AUTHORIZED -> { return false }
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.openapi.Disposable
import com.intellij.openapi.util.Disposer
import software.aws.toolkits.core.TokenConnectionSettings
import software.aws.toolkits.core.credentials.ToolkitBearerTokenProvider
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.InteractiveBearerTokenProvider
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.ProfileSdkTokenProviderWrapper
import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider

class ManagedBearerSsoConnection(
    val startUrl: String,
    val region: String,
    override val scopes: List<String>,
) : BearerSsoConnection, Disposable {
    override val id: String = ToolkitBearerTokenProvider.ssoIdentifier(startUrl, region)
    override val label: String = ToolkitBearerTokenProvider.ssoDisplayName(startUrl)

    private val provider =
        tokenConnection(
            InteractiveBearerTokenProvider(
                startUrl,
                region,
                scopes
            ),
            region
        )

    override fun getConnectionSettings(): TokenConnectionSettings = provider

    override fun dispose() {
        disposeProviderIfRequired(provider)
    }
}

class DiskSsoSessionConnection(
    val sessionProfileName: String,
    val region: String,
    displayNameOverride: String? = null
) : AwsBearerTokenConnection, Disposable {
    override val id = ToolkitBearerTokenProvider.diskSessionIdentifier(sessionProfileName)
    override val label = displayNameOverride ?: ToolkitBearerTokenProvider.diskSessionDisplayName(sessionProfileName)

    private val provider =
        tokenConnection(
            ProfileSdkTokenProviderWrapper(
                sessionName = sessionProfileName,
                region = region
            ),
            region
        )

    override fun getConnectionSettings(): TokenConnectionSettings = provider

    override fun dispose() {
        disposeProviderIfRequired(provider)
    }
}

private fun tokenConnection(provider: BearerTokenProvider, region: String) =
    TokenConnectionSettings(
        ToolkitBearerTokenProvider(provider),
        AwsRegionProvider.getInstance().get(region) ?: error("Partition data is missing for $region")
    )

private fun disposeProviderIfRequired(settings: TokenConnectionSettings) {
    val delegate = settings.tokenProvider.delegate
    if (delegate is Disposable) {
        Disposer.dispose(delegate)
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.intellij.util.messages.Topic
import java.util.EventListener

// TODO: unify with [ConnectionSettingsStateChangeNotifier]
interface ToolkitConnectionManagerListener : EventListener {
    fun activeConnectionChanged(newConnection: ToolkitConnection?)

    companion object {
        @Topic.AppLevel
        val TOPIC = Topic.create("ToolkitConnectionManagerListener active connection change", ToolkitConnectionManagerListener::class.java)
    }
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials

import com.fasterxml.jackson.core.JsonProcessingException
import com.fasterxml.jackson.databind.DeserializationFeature
import com.fasterxml.jackson.databind.MapperFeature
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule
import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
import com.fasterxml.jackson.module.kotlin.readValue
import com.intellij.execution.configurations.CommandLineTokenizer
import com.intellij.execution.configurations.GeneralCommandLine
import com.intellij.execution.process.ProcessOutput
import com.intellij.execution.util.ExecUtil
import com.intellij.openapi.util.SystemInfo
import com.intellij.openapi.util.registry.Registry
import com.intellij.util.execution.ParametersListUtil
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials
import software.amazon.awssdk.auth.credentials.AwsCredentials
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider
import software.amazon.awssdk.auth.credentials.AwsSessionCredentials
import software.amazon.awssdk.utils.SdkAutoCloseable
import software.amazon.awssdk.utils.cache.CachedSupplier
import software.amazon.awssdk.utils.cache.RefreshResult
import software.aws.toolkits.resources.message
import java.time.Instant
import java.util.Enumeration

/**
 * Similar to the SDKs ProcessCredentialsProvider, but ties in the env var system of the IDE such as getting $PATH
 */
class ToolkitCredentialProcessProvider internal constructor(
    private val command: String,
    private val parser: CredentialProcessOutputParser
) : AwsCredentialsProvider, SdkAutoCloseable {
    constructor(command: String) : this(command, DefaultCredentialProcessOutputParser)

    private val entrypoint by lazy {
        ParametersListUtil.parse(command).first()
    }
    private val cmd by lazy {
        if (SystemInfo.isWindows) {
            @Suppress("UNCHECKED_CAST")
            GeneralCommandLine("cmd", "/C", *(CommandLineTokenizer(command) as Enumeration<String>).toList().toTypedArray())
        } else {
            GeneralCommandLine("sh", "-c", command)
        }
    }
    private val processCredentialCache = CachedSupplier.builder { refresh() }.build()

    override fun resolveCredentials(): AwsCredentials = processCredentialCache.get()

    private fun refresh(): RefreshResult<AwsCredentials> {
        val timeout = Registry.intValue("aws.credentialProcess.timeout", DEFAULT_TIMEOUT)
        val output = ExecUtil.execAndGetOutput(cmd, timeout)

        if (output.isTimeout) {
            handleException(message("credentials.profile.credential_process.timeout_exception_prefix", entrypoint), output)
        }

        if (output.exitCode != 0) {
            handleException(message("credentials.profile.credential_process.execution_exception_prefix", entrypoint), output)
        }

        val result = try {
            parser.parse(output.stdout)
        } catch (e: Exception) {
            handleException(message("credentials.profile.credential_process.parse_exception_prefix"), output)
        }
        val credentials = when (val token = result.sessionToken) {
            null -> AwsBasicCredentials.create(result.accessKeyId, result.secretAccessKey)
            else -> AwsSessionCredentials.create(result.accessKeyId, result.secretAccessKey, token)
        }
        return RefreshResult.builder(credentials).staleTime(result.expiration ?: Instant.MAX).build()
    }

    private fun handleException(msgPrefix: String, process: ProcessOutput): Nothing {
        val errorOutput = process.stderr.takeIf { it.isNotBlank() }
        val msg = "$msgPrefix${errorOutput?.let { ": $it" } ?: ""}"
        throw RuntimeException(msg)
    }

    override fun close() {
        processCredentialCache.close()
    }

    private companion object {
        private const val DEFAULT_TIMEOUT = 30000
    }
}

internal data class CredentialProcessOutput(val accessKeyId: String, val secretAccessKey: String, val sessionToken: String?, val expiration: Instant?)

internal abstract class CredentialProcessOutputParser {
    abstract fun parse(input: String): CredentialProcessOutput
}

internal object DefaultCredentialProcessOutputParser : CredentialProcessOutputParser() {
    private val mapper = jacksonObjectMapper()
        .enable(MapperFeature.ACCEPT_CASE_INSENSITIVE_PROPERTIES)
        .disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
        .registerModule(JavaTimeModule())

    override fun parse(input: String): CredentialProcessOutput = try {
        mapper.readValue(input)
    } catch (e: JsonProcessingException) {
        e.clearLocation()
        throw e
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials.actions

import com.intellij.icons.AllIcons
import com.intellij.ide.BrowserUtil
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.project.DumbAwareAction
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
import java.net.URI

class CredentialsHelpAction : DumbAwareAction(AllIcons.General.ContextHelp) {
    override fun actionPerformed(e: AnActionEvent) {
        // TODO: update
        runInEdt {
            BrowserUtil.browse(URI(CodeWhispererConstants.CODEWHISPERER_LOGIN_HELP_URI))
        }
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials.actions

import com.intellij.icons.AllIcons
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.project.DumbAwareAction
import software.aws.toolkits.jetbrains.core.credentials.ToolkitAddConnectionDialog
import software.aws.toolkits.telemetry.UiTelemetry

class ExplorerNewConnectionAction : DumbAwareAction(AllIcons.General.Add) {
    override fun displayTextInToolbar() = true

    override fun actionPerformed(e: AnActionEvent) {
        e.project?.let {
            runInEdt {
                ToolkitAddConnectionDialog(it).showAndGet()
                UiTelemetry.click(it, "devtools_connectToAws")
            }
        }
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials.actions

import com.intellij.icons.AllIcons
import com.intellij.openapi.actionSystem.ActionGroup
import com.intellij.openapi.actionSystem.ActionManager
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.actionSystem.DefaultActionGroup
import com.intellij.openapi.project.DumbAwareAction
import com.intellij.openapi.ui.popup.JBPopupFactory
import software.aws.toolkits.jetbrains.core.credentials.ConnectionSettingsMenuBuilder.Companion.connectionSettingsMenuBuilder
import software.aws.toolkits.jetbrains.core.credentials.lazyGetUnauthedBearerConnections
import software.aws.toolkits.resources.message

class MoreConnectionActionsAction : DumbAwareAction(AllIcons.Actions.MoreHorizontal) {
    override fun update(e: AnActionEvent) {
        e.presentation.icon = if (lazyGetUnauthedBearerConnections().isEmpty()) AllIcons.Actions.MoreHorizontal else AllIcons.General.Warning
    }

    override fun actionPerformed(e: AnActionEvent) {
        val baseActions = ActionManager.getInstance().getAction("aws.toolkit.toolwindow.credentials.rightGroup.more.group") as ActionGroup
        val identityActions = connectionSettingsMenuBuilder().apply { withIndividualIdentityActions(e.project) }.build()

        JBPopupFactory.getInstance().createActionGroupPopup(
            message("settings.title"),
            DefaultActionGroup(identityActions, baseActions),
            e.dataContext,
            JBPopupFactory.ActionSelectionAid.SPEEDSEARCH,
            true
        ).showInBestPositionFor(e.dataContext)
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials.actions

import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.project.DumbAwareAction
import software.aws.toolkits.jetbrains.core.credentials.ToolkitAddConnectionDialog

class NewConnectionAction : DumbAwareAction() {
    override fun actionPerformed(e: AnActionEvent) {
        e.project?.let {
            runInEdt {
                ToolkitAddConnectionDialog(it).showAndGet()
            }
        }
    }
}


// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials.metadataservice

import com.intellij.openapi.extensions.ExtensionNotApplicableException
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider
import software.amazon.awssdk.auth.credentials.ContainerCredentialsProvider
import software.amazon.awssdk.core.SdkSystemSetting
import software.aws.toolkits.core.credentials.CredentialIdentifier
import software.aws.toolkits.core.credentials.CredentialProviderFactory
import software.aws.toolkits.core.credentials.CredentialSourceId
import software.aws.toolkits.core.credentials.CredentialType
import software.aws.toolkits.core.credentials.CredentialsChangeEvent
import software.aws.toolkits.core.credentials.CredentialsChangeListener
import software.aws.toolkits.core.region.AwsRegion
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.jetbrains.services.caws.CawsConstants

class ContainerCredentialProviderFactory : CredentialProviderFactory {
    init {
        if (System.getenv(CawsConstants.CAWS_ENV_ID_VAR) != null) {
            throw ExtensionNotApplicableException.INSTANCE
        }
    }

    override val id = "ContainerCredentialProviderFactory"
    override val credentialSourceId: CredentialSourceId = CredentialSourceId.Ecs

    private val containerCredIdentifier by lazy {
        credentialId()
    }

    override fun setUp(credentialLoadCallback: CredentialsChangeListener) {
        // deviates from SDK behavior by treating the empty value as unset
        val credSettings = arrayOf(SdkSystemSetting.AWS_CONTAINER_CREDENTIALS_RELATIVE_URI, SdkSystemSetting.AWS_CONTAINER_CREDENTIALS_FULL_URI)
        if (credSettings.none { it.stringValue.orElse("").isNotBlank() }) {
            getLogger<ContainerCredentialProviderFactory>().debug {
                "Skipping container credential provider since container credentials environment variables were not available"
            }

            return
        }

        credentialLoadCallback(
            CredentialsChangeEvent(
                added = listOf(containerCredIdentifier),
                modified = emptyList(),
                removed = emptyList()
            )
        )
    }

    override fun createAwsCredentialProvider(providerId: CredentialIdentifier, region: AwsRegion): AwsCredentialsProvider =
        ContainerCredentialsProvider.builder()
            .asyncCredentialUpdateEnabled(false)
            .build()

    companion object {
        const val FACTORY_ID = "ContainerCredentialProviderFactory"

        fun credentialId() = object : CredentialIdentifier {
            override val id: String = "containerRoleCredential"
            override val displayName = "ecs:containerRole"
            override val factoryId = FACTORY_ID
            override val credentialType = CredentialType.EcsMetadata
            override val defaultRegionId = System.getenv(SdkSystemSetting.AWS_REGION.environmentVariable())
        }
    }
}


// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials.metadataservice

import com.intellij.openapi.extensions.ExtensionNotApplicableException
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider
import software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider
import software.amazon.awssdk.core.SdkSystemSetting
import software.amazon.awssdk.regions.internal.util.EC2MetadataUtils
import software.aws.toolkits.core.credentials.CredentialIdentifier
import software.aws.toolkits.core.credentials.CredentialProviderFactory
import software.aws.toolkits.core.credentials.CredentialSourceId
import software.aws.toolkits.core.credentials.CredentialType
import software.aws.toolkits.core.credentials.CredentialsChangeEvent
import software.aws.toolkits.core.credentials.CredentialsChangeListener
import software.aws.toolkits.core.region.AwsRegion
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.warn
import software.aws.toolkits.jetbrains.services.caws.CawsConstants

class InstanceRoleCredentialProviderFactory : CredentialProviderFactory {
    init {
        if (System.getenv(CawsConstants.CAWS_ENV_ID_VAR) != null) {
            throw ExtensionNotApplicableException.INSTANCE
        }
    }

    override val id = FACTORY_ID
    override val credentialSourceId: CredentialSourceId = CredentialSourceId.Ec2

    private val instanceRoleCredIdentifier: CredentialIdentifier by lazy {
        object : CredentialIdentifier {
            override val id: String = "ec2InstanceRoleCredential"
            override val displayName = "ec2:instanceProfile"
            override val factoryId = FACTORY_ID
            override val credentialType = CredentialType.Ec2Metadata
            override val defaultRegionId = try {
                EC2MetadataUtils.getEC2InstanceRegion()
            } catch (e: Exception) {
                LOG.warn(e) { "Failed to query instance region from ec2 instance metadata" }
                null
            }
        }
    }

    private val provider = InstanceProfileCredentialsProvider.builder()
        .asyncCredentialUpdateEnabled(false)
        .build()

    override fun setUp(credentialLoadCallback: CredentialsChangeListener) {
        if (SdkSystemSetting.AWS_EC2_METADATA_DISABLED.booleanValue.orElse(false)) {
            LOG.debug { "EC2 metadata provider disabled by system setting" }
            return
        }

        val endpoint = SdkSystemSetting.AWS_EC2_METADATA_SERVICE_ENDPOINT.stringValue.orElse("")
        if (endpoint.isBlank()) {
            LOG.debug { "Skipping instance role credential provider since endpoint was blank" }
            return
        }

        try {
            provider.resolveCredentials()
        } catch (e: Exception) {
            provider.close()
            LOG.debug { "Instance role credential provider failed to resolve credentials" }
            return
        }

        credentialLoadCallback(
            CredentialsChangeEvent(
                added = listOf(instanceRoleCredIdentifier),
                modified = emptyList(),
                removed = emptyList()
            )
        )
    }

    override fun createAwsCredentialProvider(providerId: CredentialIdentifier, region: AwsRegion): AwsCredentialsProvider =
        provider

    companion object {
        const val FACTORY_ID = "InstanceRoleCredentialProviderFactory"
        private val LOG = getLogger<InstanceRoleCredentialProviderFactory>()
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials.pinning

import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
import software.aws.toolkits.jetbrains.core.credentials.BearerSsoConnection
import software.aws.toolkits.jetbrains.core.credentials.ManagedBearerSsoConnection
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
import software.aws.toolkits.jetbrains.core.credentials.sono.CODECATALYST_SCOPES
import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_URL

class CodeCatalystConnection : FeatureWithPinnedConnection {
    override val featureId: String = "aws.codecatalyst"
    override val featureName: String = "CodeCatalyst"
    override fun supportsConnectionType(connection: ToolkitConnection): Boolean {
        if (connection is AwsBearerTokenConnection) {
            if (connection is ManagedBearerSsoConnection && connection.startUrl != SONO_URL) {
                LOG.debug { "Rejecting ${connection.id} since it's not a AWS Builder ID connection" }
                // doesn't support arbitrary SSO
                return false
            }

            if (connection is BearerSsoConnection && !CODECATALYST_SCOPES.all { it in connection.scopes }) {
                LOG.debug { "Rejecting ${connection.id} since it's missing a required scope" }
                return false
            }

            return true
        }

        // only supports bearer connections
        return false
    }

    companion object {
        private val LOG = getLogger<CodeCatalystConnection>()
        fun getInstance() = FeatureWithPinnedConnection.EP_NAME.findExtensionOrFail(CodeCatalystConnection::class.java)
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials.pinning

import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
import software.aws.toolkits.jetbrains.core.credentials.BearerSsoConnection
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
import software.aws.toolkits.jetbrains.core.credentials.sono.CODEWHISPERER_SCOPES

class CodeWhispererConnection : FeatureWithPinnedConnection {
    override val featureId = "aws.codewhisperer"
    override val featureName = "CodeWhisperer"

    override fun supportsConnectionType(connection: ToolkitConnection): Boolean {
        if (connection is AwsBearerTokenConnection) {
            if (connection is BearerSsoConnection) {
                return CODEWHISPERER_SCOPES.all { it in connection.scopes }
            }

            return true
        }

        return false
    }

    companion object {
        fun getInstance() = FeatureWithPinnedConnection.EP_NAME.findExtensionOrFail(CodeWhispererConnection::class.java)
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials.pinning

import com.intellij.icons.AllIcons
import com.intellij.openapi.Disposable
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.application.ModalityState
import com.intellij.openapi.components.PersistentStateComponent
import com.intellij.openapi.components.State
import com.intellij.openapi.components.Storage
import com.intellij.openapi.components.service
import com.intellij.openapi.extensions.ExtensionPointName
import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.DialogWrapper
import com.intellij.openapi.ui.MessageDialogBuilder
import org.jetbrains.annotations.VisibleForTesting
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
import software.aws.toolkits.jetbrains.core.credentials.ToolkitAuthManager
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
import software.aws.toolkits.jetbrains.core.help.HelpIds
import software.aws.toolkits.jetbrains.utils.computeOnEdt
import software.aws.toolkits.jetbrains.utils.notifyInfo
import software.aws.toolkits.resources.message
import software.aws.toolkits.telemetry.UiTelemetry
import java.util.concurrent.ConcurrentHashMap

interface FeatureWithPinnedConnection {
    val featureId: String
    val featureName: String

    fun supportsConnectionType(connection: ToolkitConnection): Boolean

    companion object {
        val EP_NAME = ExtensionPointName<FeatureWithPinnedConnection>("aws.toolkit.connection.pinned.feature")
    }
}

interface ConnectionPinningManager {
    fun isFeaturePinned(feature: FeatureWithPinnedConnection): Boolean
    fun getPinnedConnection(feature: FeatureWithPinnedConnection): ToolkitConnection?
    fun setPinnedConnection(feature: FeatureWithPinnedConnection, newConnection: ToolkitConnection?)

    fun maybePinFeatures(oldConnection: ToolkitConnection?, newConnection: ToolkitConnection, features: List<FeatureWithPinnedConnection>)

    companion object {
        fun getInstance(): ConnectionPinningManager = service()
    }
}

@State(name = "connectionPinningManager", storages = [Storage("aws.xml")])
class DefaultConnectionPinningManager :
    ConnectionPinningManager,
    PersistentStateComponent<ConnectionPinningManagerState>,
    Disposable {
    @VisibleForTesting
    internal var shouldPinConnections: Boolean? = null

    private val pinnedConnections = ConcurrentHashMap<String, ToolkitConnection>()

    init {
        ApplicationManager.getApplication().messageBus.connect(this).subscribe(
            BearerTokenProviderListener.TOPIC,
            object : BearerTokenProviderListener {
                override fun invalidate(providerId: String) {
                    pinnedConnections.entries.removeIf { (_, v) -> v.id == providerId }
                }
            }
        )
    }

    override fun isFeaturePinned(feature: FeatureWithPinnedConnection) = getPinnedConnection(feature) != null

    override fun getPinnedConnection(feature: FeatureWithPinnedConnection): ToolkitConnection? =
        pinnedConnections[feature.featureId].let { connection ->
            if (connection == null || ToolkitAuthManager.getInstance().getConnection(connection.id) == null) {
                null
            } else {
                connection
            }
        }

    override fun setPinnedConnection(feature: FeatureWithPinnedConnection, newConnection: ToolkitConnection?) {
        if (newConnection == null) {
            pinnedConnections.remove(feature.featureId)
        } else {
            pinnedConnections[feature.featureId] = newConnection
        }

        ApplicationManager.getApplication().messageBus.syncPublisher(ConnectionPinningManagerListener.TOPIC).pinnedConnectionChanged(feature, newConnection)
    }

    override fun maybePinFeatures(oldConnection: ToolkitConnection?, newConnection: ToolkitConnection, features: List<FeatureWithPinnedConnection>) {
        val featuresString = if (features.size == 1) {
            features.first().featureName
        } else {
            "${features.dropLast(1).joinToString(",") { it.featureName }} and ${features.last().featureName}"
        }

        var connectionToPin = if (oldConnection is AwsBearerTokenConnection) oldConnection else newConnection
        if (showDialogIfNeeded(oldConnection, newConnection, featuresString)) {
            features.forEach {
                setPinnedConnection(it, connectionToPin)
            }
            notifyInfo(message("credentials.switch.notification.title", featuresString, connectionToPin.label))
        }
    }

    override fun getState() = ConnectionPinningManagerState(
        shouldPinConnections,
        pinnedConnections.entries.associate { (k, v) -> k to v.id }
    )

    override fun loadState(state: ConnectionPinningManagerState) {
        val authManager = ToolkitAuthManager.getInstance()

        shouldPinConnections = state.shouldPinConnections

        pinnedConnections.clear()
        pinnedConnections.putAll(
            state.pinnedConnections.entries.mapNotNull { (k, v) ->
                authManager.getConnection(v)?.let { k to it }
            }
        )
    }

    override fun dispose() {}

    @VisibleForTesting
    internal fun showDialogIfNeeded(oldConnection: ToolkitConnection?, newConnection: ToolkitConnection, featuresString: String, project: Project? = null) =
        shouldPinConnections.let { shouldPinConnections ->
            if (shouldPinConnections == null) {
                val bearerTokenConnectionName = bearerTokenConnectionString(oldConnection, newConnection)

                computeOnEdt(ModalityState.defaultModalityState()) {
                    MessageDialogBuilder.yesNo(
                        message("credentials.switch.confirmation.title", featuresString, bearerTokenConnectionName),
                        message("credentials.switch.confirmation.comment", featuresString, bearerTokenConnectionName, message("iam.name"))
                    )
                        .yesText(message("credentials.switch.confirmation.yes"))
                        .noText(message("credentials.switch.confirmation.no"))
                        .doNotAsk(object : com.intellij.openapi.ui.DoNotAskOption.Adapter() {
                            override fun rememberChoice(isSelected: Boolean, exitCode: Int) {
                                if (isSelected) {
                                    this@DefaultConnectionPinningManager.shouldPinConnections = exitCode == DialogWrapper.OK_EXIT_CODE
                                }
                            }
                        })
                        .icon(AllIcons.General.QuestionDialog)
                        .help(HelpIds.EXPLORER_CREDS_HELP.id)
                        .ask(project).apply {
                            if (this) {
                                UiTelemetry.click(project, "connection_multiple_auths_yes")
                            } else {
                                UiTelemetry.click(project, "connection_multiple_auths_no")
                            }
                        }
                }
            } else {
                shouldPinConnections
            }
        }

    private fun bearerTokenConnectionString(oldConnection: ToolkitConnection?, newConnection: ToolkitConnection): String {
        val connection = if (oldConnection is AwsBearerTokenConnection) oldConnection else newConnection
        return if (connection.isSono()) message("aws_builder_id.service_name") else message("iam_identity_center.name")
    }
}

data class ConnectionPinningManagerState(
    var shouldPinConnections: Boolean? = null,
    var pinnedConnections: Map<String, String> = emptyMap()
)


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials.pinning

import com.intellij.util.messages.Topic
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
import java.util.EventListener

interface ConnectionPinningManagerListener : EventListener {
    fun pinnedConnectionChanged(feature: FeatureWithPinnedConnection, newConnection: ToolkitConnection?)

    companion object {
        @Topic.AppLevel
        val TOPIC = Topic.create("Feature pinned active connection change", ConnectionPinningManagerListener::class.java)
    }
}


// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials.profiles

enum class CredentialSourceType {
    EC2_INSTANCE_METADATA, ECS_CONTAINER, ENVIRONMENT;

    companion object {
        fun parse(value: String): CredentialSourceType {
            if (value.equals("Ec2InstanceMetadata", ignoreCase = true)) {
                return EC2_INSTANCE_METADATA
            } else if (value.equals("EcsContainer", ignoreCase = true)) {
                return ECS_CONTAINER
            } else if (value.equals("Environment", ignoreCase = true)) {
                return ENVIRONMENT
            }
            throw IllegalArgumentException("'$value' is not a valid credential_source")
        }
    }
}


// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials.profiles

import software.amazon.awssdk.core.SdkSystemSetting
import software.amazon.awssdk.profiles.Profile
import software.amazon.awssdk.profiles.ProfileProperty

/**
 * Retrieves the EC2 metadata endpoint based on profile file, env var, and Java system properties
 *
 * https://github.com/aws/aws-sdk-java-v2/blob/5fb447594313ab1ab9b9c0ead0ed7cb906b06e93/core/auth/src/test/java/software/amazon/awssdk/auth/credentials/internal/Ec2MetadataConfigProviderEndpointResolutionTest.java
 */
object Ec2MetadataConfigProvider {
    /**
     * Default IPv4 endpoint for the Amazon EC2 Instance Metadata Service.
     */
    private const val EC2_METADATA_SERVICE_URL_IPV4 = "http://169.254.169.254"

    /**
     * Default IPv6 endpoint for the Amazon EC2 Instance Metadata Service.
     */
    private const val EC2_METADATA_SERVICE_URL_IPV6 = "http://[fd00:ec2::254]"

    private enum class EndpointMode {
        IPV4, IPV6;

        companion object {
            fun fromValue(s: String?): EndpointMode = s?.let { _ ->
                values().find { it.name.equals(s, ignoreCase = true) }
            } ?: throw IllegalArgumentException("Unrecognized value for endpoint mode: '$s'")
        }
    }

    fun Profile.getEc2MedataEndpoint(): String = this.getEndpointOverride() ?: when (this.getEndpointMode()) {
        EndpointMode.IPV4 -> EC2_METADATA_SERVICE_URL_IPV4
        EndpointMode.IPV6 -> EC2_METADATA_SERVICE_URL_IPV6
    }

    private fun Profile.getEndpointMode(): EndpointMode {
        val endpointMode = SdkSystemSetting.AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE.nonDefaultStringValue
        val endpointModeString = if (endpointMode.isPresent) {
            endpointMode.get()
        } else {
            configFileEndpointMode(this) ?: SdkSystemSetting.AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE.defaultValue()
        }
        return EndpointMode.fromValue(endpointModeString)
    }

    private fun Profile.getEndpointOverride(): String? {
        val endpointOverride = SdkSystemSetting.AWS_EC2_METADATA_SERVICE_ENDPOINT.nonDefaultStringValue
        return if (endpointOverride.isPresent) {
            endpointOverride.get()
        } else {
            configFileEndpointOverride(this)
        }
    }

    private fun configFileEndpointMode(profile: Profile): String? = profile.property(ProfileProperty.EC2_METADATA_SERVICE_ENDPOINT_MODE).orElse(null)

    private fun configFileEndpointOverride(profile: Profile): String? = profile.property(ProfileProperty.EC2_METADATA_SERVICE_ENDPOINT).orElse(null)
}


// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.core.credentials.profiles

import org.jetbrains.annotations.TestOnly
import software.amazon.awssdk.auth.credentials.AwsCredentials
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider
import software.amazon.awssdk.profiles.Profile
import software.amazon.awssdk.profiles.ProfileProperty
import software.amazon.awssdk.regions.Region
import software.amazon.awssdk.services.sts.StsClient
import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider
import software.amazon.awssdk.services.sts.model.AssumeRoleRequest
import software.amazon.awssdk.utils.SdkAutoCloseable
import software.aws.toolkits.core.region.AwsRegion
import software.aws.toolkits.jetbrains.core.AwsClientManager
import software.aws.toolkits.jetbrains.core.credentials.promptForMfaToken
import java.util.function.Supplier

class ProfileAssumeRoleProvider(@get:TestOnly internal val parentProvider: AwsCredentialsProvider, region: AwsRegion, profile: Profile) :
    AwsCredentialsProvider, SdkAutoCloseable {
    private val stsClient: StsClient
    private val credentialsProvider: StsAssumeRoleCredentialsProvider

    init {
        val roleArn = profile.requiredProperty(ProfileProperty.ROLE_ARN)
        val roleSessionName = profile.property(ProfileProperty.ROLE_SESSION_NAME).orElseGet { "aws-toolkit-jetbrains-${System.currentTimeMillis()}" }
        val externalId = profile.property(ProfileProperty.EXTERNAL_ID).orElse(null)
        val mfaSerial = profile.property(ProfileProperty.MFA_SERIAL).orElse(null)

        // https://docs.aws.amazon.com/sdkref/latest/guide/setting-global-duration_seconds.html
        val durationSecs = profile.property(ProfileProperty.DURATION_SECONDS).map { it.toIntOrNull() }.orElse(null) ?: 3600

        stsClient = AwsClientManager.getInstance().createUnmanagedClient(parentProvider, Region.of(region.id))

        credentialsProvider = StsAssumeRoleCredentialsProvider.builder()
            .stsClient(stsClient)
            .refreshRequest(
                Supplier {
                    createAssumeRoleRequest(
                        profile.name(),
                        mfaSerial,
                        roleArn,
                        roleSessionName,
                        externalId,
                        durationSecs
                    )
                }
            )
            .build()
    }

    private fun createAssumeRoleRequest(
        profileName: String,
        mfaSerial: String?,
        roleArn: String,
        roleSessionName: String?,
        externalId: String?,
        durationSeconds: Int
    ): AssumeRoleRequest {
        val requestBuilder = AssumeRoleRequest.builder()
            .roleArn(roleArn)
            .roleSessionName(roleSessionName)
            .externalId(externalId)
            .durationSeconds(durationSeconds)

        mfaSerial?.let { _ ->
            requestBuilder
                .serialNumber(mfaSerial)
                .tokenCode(promptForMfaToken(profileName, mfaSerial))
        }

        return requestBuilder.build()
    }

    override fun resolveCredentials(): AwsCredentials = credentialsProvider.resolveCredentials()

    override fun close() {
        credentialsProvider.close()
        (parentProvider as? SdkAutoCloseable)?.close()
        stsClient.close()
    }
}