diff --git a/quickstart/201-vm-disk-encryption-extension/main.tf b/quickstart/201-vm-disk-encryption-extension/main.tf
index 7069557d..a112993d 100644
--- a/quickstart/201-vm-disk-encryption-extension/main.tf
+++ b/quickstart/201-vm-disk-encryption-extension/main.tf
@@ -1,13 +1,17 @@
 resource "azurerm_resource_group" "example" {
-  name     = "${var.name_prefix}-rg"
+  name     = "${random_pet.prefix.id}-rg"
   location = var.location
 }
 
 // Key Vault Key
 data "azurerm_client_config" "current" {}
 
+locals {
+  current_user_object_id = coalesce(var.msi_id, data.azurerm_client_config.current.object_id)
+}
+
 resource "azurerm_key_vault" "example" {
-  name                        = "${var.name_prefix}-kv"
+  name                        = "${random_pet.prefix.id}-kv"
   location                    = azurerm_resource_group.example.location
   resource_group_name         = azurerm_resource_group.example.name
   tenant_id                   = data.azurerm_client_config.current.tenant_id
@@ -20,7 +24,7 @@ resource "azurerm_key_vault" "example" {
 resource "azurerm_key_vault_access_policy" "service-principal" {
   key_vault_id = azurerm_key_vault.example.id
   tenant_id    = data.azurerm_client_config.current.tenant_id
-  object_id    = data.azurerm_client_config.current.object_id
+  object_id    = local.current_user_object_id
 
   key_permissions = [
     "Create",
@@ -58,21 +62,21 @@ resource "azurerm_key_vault_key" "example" {
 
 // Virtual Machine
 resource "azurerm_virtual_network" "example" {
-  name                = "${var.name_prefix}-vnet"
+  name                = "${random_pet.prefix.id}-vnet"
   address_space       = ["10.0.0.0/16"]
   location            = azurerm_resource_group.example.location
   resource_group_name = azurerm_resource_group.example.name
 }
 
 resource "azurerm_subnet" "example" {
-  name                 = "${var.name_prefix}-subnet"
+  name                 = "${random_pet.prefix.id}-subnet"
   resource_group_name  = azurerm_resource_group.example.name
   virtual_network_name = azurerm_virtual_network.example.name
   address_prefixes     = ["10.0.2.0/24"]
 }
 
 resource "azurerm_network_interface" "example" {
-  name                = "${var.name_prefix}-nic"
+  name                = "${random_pet.prefix.id}-nic"
   location            = azurerm_resource_group.example.location
   resource_group_name = azurerm_resource_group.example.name
 
@@ -83,8 +87,8 @@ resource "azurerm_network_interface" "example" {
   }
 }
 
-resource "azurerm_linux_virtual_machine" "example" {
-  name                = "${var.name_prefix}-vm"
+resource "azurerm_linux_virtual_machine" "main" {
+  name                = "${random_pet.prefix.id}-vm"
   resource_group_name = azurerm_resource_group.example.name
   location            = azurerm_resource_group.example.location
   size                = "Standard_D2s_v3"
@@ -95,7 +99,7 @@ resource "azurerm_linux_virtual_machine" "example" {
 
   admin_ssh_key {
     username   = "azureuser"
-    public_key = var.vm_public_key
+    public_key = tls_private_key.vm_key.public_key_openssh
   }
 
   source_image_reference {
@@ -118,7 +122,7 @@ resource "azurerm_virtual_machine_extension" "example" {
   type                       = "AzureDiskEncryptionForLinux"
   type_handler_version       = "1.1"
   auto_upgrade_minor_version = false
-  virtual_machine_id         = azurerm_linux_virtual_machine.example.id
+  virtual_machine_id         = azurerm_linux_virtual_machine.main.id
 
   settings = jsonencode({
     "EncryptionOperation"    = "EnableEncryption"
@@ -130,3 +134,18 @@ resource "azurerm_virtual_machine_extension" "example" {
     "VolumeType"             = "All"
   })
 }
+
+resource "local_sensitive_file" "lsf" {
+  content  = tls_private_key.vm_key.private_key_pem
+  filename = "key.pem"
+}
+
+resource "random_pet" "prefix" {
+  length = 1
+  prefix = var.prefix
+}
+
+resource "tls_private_key" "vm_key" {
+  algorithm = "RSA"
+  rsa_bits  = 4096
+}
\ No newline at end of file
diff --git a/quickstart/201-vm-disk-encryption-extension/outputs.tf b/quickstart/201-vm-disk-encryption-extension/outputs.tf
new file mode 100644
index 00000000..e69de29b
diff --git a/quickstart/201-vm-disk-encryption-extension/providers.tf b/quickstart/201-vm-disk-encryption-extension/providers.tf
index 684448ec..0abec932 100644
--- a/quickstart/201-vm-disk-encryption-extension/providers.tf
+++ b/quickstart/201-vm-disk-encryption-extension/providers.tf
@@ -1,10 +1,22 @@
 terraform {
-  required_version = ">=1.0"
+  required_version = ">= 1.0"
 
   required_providers {
     azurerm = {
       source  = "hashicorp/azurerm"
-      version = "~>3.0"
+      version = ">= 3.0, < 4.0"
+    }
+    local = {
+      source  = "hashicorp/local"
+      version = "2.3.0"
+    }
+    random = {
+      source  = "hashicorp/random"
+      version = ">= 3.0"
+    }
+    tls = {
+      source  = "hashicorp/tls"
+      version = "4.0.4"
     }
   }
 }
diff --git a/quickstart/201-vm-disk-encryption-extension/variables.tf b/quickstart/201-vm-disk-encryption-extension/variables.tf
index b5e9aa6e..69f3bf18 100644
--- a/quickstart/201-vm-disk-encryption-extension/variables.tf
+++ b/quickstart/201-vm-disk-encryption-extension/variables.tf
@@ -1,14 +1,17 @@
 variable "location" {
   type        = string
+  default     = "eastus"
   description = "Location where resources will be created"
 }
 
-variable "name_prefix" {
+variable "msi_id" {
   type        = string
-  description = "Prefix of the resource name"
+  default     = null
+  description = "If you're executing the test with user assigned identity, please pass the identity principal id to this variable."
 }
 
-variable "vm_public_key" {
+variable "prefix" {
   type        = string
-  description = "Public key of the Virtual Machine"
-}
+  default     = "vm-disk-e-e"
+  description = "Prefix of the resource name"
+}
\ No newline at end of file