diff --git a/README.md b/README.md
index 98914af..70ca540 100644
--- a/README.md
+++ b/README.md
@@ -178,3 +178,12 @@ No modules.
 | <a name="output_mwaa_status"></a> [mwaa\_status](#output\_mwaa\_status) | The status of the Amazon MWAA Environment |
 | <a name="output_mwaa_webserver_url"></a> [mwaa\_webserver\_url](#output\_mwaa\_webserver\_url) | The webserver URL of the MWAA Environment |
 <!-- END OF PRE-COMMIT-TERRAFORM DOCS HOOK -->
+
+## Troubleshooting
+Tools and sample code provided by AWS Premium Support for troubleshooting your MWAA Deployments : 
+
+```hcl
+$ pip3 install boto3 --upgrade --user
+$ git clone https://github.com/awslabs/aws-support-tools.git
+$ python3 aws-support-tools/MWAA/verify_env/verify_env.py --envname YOUR_ENV_NAME_HERE
+```
diff --git a/data.tf b/data.tf
index 8d0e323..c652d2a 100644
--- a/data.tf
+++ b/data.tf
@@ -4,6 +4,13 @@ data "aws_region" "current" {}
 
 data "aws_caller_identity" "current" {}
 
+data "aws_s3_bucket" "selected" {
+  bucket = var.source_bucket_name
+}
+
+data "aws_kms_key" "by_alias" {
+  key_id = var.kms_key
+}
 # ---------------------------------------------------------------------------------------------------------------------
 # MWAA Role
 # ---------------------------------------------------------------------------------------------------------------------
@@ -21,23 +28,6 @@ data "aws_iam_policy_document" "mwaa_assume" {
       identifiers = ["airflow-env.amazonaws.com"]
     }
 
-    principals {
-      type        = "Service"
-      identifiers = ["batch.amazonaws.com"]
-    }
-
-    principals {
-      type        = "Service"
-      identifiers = ["ssm.amazonaws.com"]
-    }
-    principals {
-      type        = "Service"
-      identifiers = ["lambda.amazonaws.com"]
-    }
-    principals {
-      type        = "Service"
-      identifiers = ["s3.amazonaws.com"]
-    }
   }
 }
 #tfsec:ignore:AWS099
@@ -45,8 +35,7 @@ data "aws_iam_policy_document" "mwaa" {
   statement {
     effect = "Allow"
     actions = [
-      "airflow:PublishMetrics",
-      "airflow:CreateWebLoginToken"
+      "airflow:PublishMetrics"
     ]
     resources = [
       "arn:${data.aws_partition.current.id}:airflow:${data.aws_region.current.name}:${data.aws_caller_identity.current.account_id}:environment/${var.name}"
@@ -58,11 +47,20 @@ data "aws_iam_policy_document" "mwaa" {
       "s3:*"
     ]
     resources = [
-      local.source_bucket_arn,
-      "${local.source_bucket_arn}/*"
+      data.aws_s3_bucket.selected.arn,
+      "${data.aws_s3_bucket.selected.arn}/*"
+    ]
+  }
+# Restrict Public Access 
+  statement {
+    effect = "Allow"
+    actions =  [
+        "s3:GetAccountPublicAccessBlock"
+    ]
+    resources = [
+      "*"
     ]
   }
-
   statement {
     effect = "Allow"
     actions = [
@@ -83,10 +81,7 @@ data "aws_iam_policy_document" "mwaa" {
     effect = "Allow"
     actions = [
       "logs:DescribeLogGroups",
-      "cloudwatch:PutMetricData",
-      "batch:DescribeJobs",
-      "batch:ListJobs",
-      "eks:*"
+      "cloudwatch:PutMetricData"
     ]
     resources = [
       "*"
@@ -109,57 +104,28 @@ data "aws_iam_policy_document" "mwaa" {
   }
 
   statement {
-    effect = "Allow"
-    actions = [
+    effect        = "Allow"
+    actions       = [
       "kms:Decrypt",
       "kms:DescribeKey",
       "kms:GenerateDataKey*",
-      "kms:Encrypt"
-    ]
-    not_resources = [
-      "arn:${data.aws_partition.current.id}:kms:*:${data.aws_caller_identity.current.account_id}:key/*"
-    ]
+      "kms:Encrypt",
+      "kms:PutKeyPolicy"
+    ]
+    resources     = data.aws_kms_key.by_alias.arn != null ? [
+      data.aws_kms_key.by_alias.arn
+    ] : []
+    not_resources = data.aws_kms_key.by_alias.arn == null ? [
+      "arn:aws:kms:*:${data.aws_caller_identity.current.account_id}:key/*"
+    ] : []
     condition {
       test     = "StringLike"
-      variable = "kms:ViaService"
-
-      values = [
+      values   = data.aws_kms_key.by_alias.arn != null ? [
+        "sqs.${data.aws_region.current.name}.amazonaws.com" ] : [
         "sqs.${data.aws_region.current.name}.amazonaws.com"
       ]
+      variable = "kms:ViaService"
     }
   }
 
-  statement {
-    effect = "Allow"
-    actions = [
-      "batch:*",
-    ]
-    resources = [
-      "arn:${data.aws_partition.current.id}:batch:*:${data.aws_caller_identity.current.account_id}:*"
-    ]
-  }
-
-  statement {
-    effect = "Allow"
-    actions = [
-      "ssm:*"
-    ]
-    resources = [
-      "arn:${data.aws_partition.current.id}:ssm:${data.aws_region.current.name}:${data.aws_caller_identity.current.account_id}:parameter/*"
-    ]
-  }
-
-  statement {
-    effect = "Allow"
-    actions = [
-      "logs:*"
-    ]
-    resources = ["arn:${data.aws_partition.current.id}:logs:${data.aws_region.current.name}:${data.aws_caller_identity.current.account_id}:log-group:/aws/lambda/*"]
-  }
-
-  statement {
-    effect    = "Allow"
-    actions   = ["cloudwatch:*"]
-    resources = ["arn:${data.aws_partition.current.id}:logs:${data.aws_region.current.name}:${data.aws_caller_identity.current.account_id}:log-group:/aws/lambda/*"]
-  }
 }
diff --git a/locals.tf b/locals.tf
index a741a93..c82f9ec 100644
--- a/locals.tf
+++ b/locals.tf
@@ -3,7 +3,7 @@ locals {
 
   security_group_ids = var.create_security_group ? concat([aws_security_group.mwaa[0].id], var.security_group_ids) : var.security_group_ids
 
-  source_bucket_arn = var.create_s3_bucket ? aws_s3_bucket.mwaa[0].arn : var.source_bucket_arn
+  source_bucket_arn = var.create_s3_bucket ? aws_s3_bucket.mwaa[0].arn : data.aws_s3_bucket.selected.arn
 
   default_airflow_configuration_options = {
     "logging.logging_level" = "INFO"
diff --git a/main.tf b/main.tf
index 1d650b5..e6b3e0a 100644
--- a/main.tf
+++ b/main.tf
@@ -7,7 +7,7 @@ resource "aws_mwaa_environment" "mwaa" {
   environment_class = var.environment_class
   min_workers       = var.min_workers
   max_workers       = var.max_workers
-  kms_key           = var.kms_key
+  kms_key           = data.aws_kms_key.by_alias.arn
 
   dag_s3_path                      = var.dag_s3_path
   plugins_s3_object_version        = var.plugins_s3_object_version
@@ -20,7 +20,7 @@ resource "aws_mwaa_environment" "mwaa" {
   execution_role_arn               = local.execution_role_arn
   airflow_configuration_options    = local.airflow_configuration_options
 
-  source_bucket_arn               = local.source_bucket_arn
+  source_bucket_arn               = data.aws_s3_bucket.selected.arn
   webserver_access_mode           = var.webserver_access_mode
   weekly_maintenance_window_start = var.weekly_maintenance_window_start
 
diff --git a/variables.tf b/variables.tf
index 2d3ac3f..5ce6d6d 100644
--- a/variables.tf
+++ b/variables.tf
@@ -206,11 +206,11 @@ variable "source_bucket_name" {
   default     = null
 }
 
-variable "source_bucket_arn" {
-  description = "(Required) The Amazon Resource Name (ARN) of your Amazon S3 storage bucket. For example, arn:aws:s3:::airflow-mybucketname"
-  type        = string
-  default     = null
-}
+#variable "source_bucket_arn" {
+#  description = "(Required) The Amazon Resource Name (ARN) of your Amazon S3 storage bucket. For example, arn:aws:s3:::airflow-mybucketname"
+#  type        = string
+#  default     = null
+#}
 
 #----------------------------------------------------------------
 # MWAA Security groups