azure vm-custom-script-extension: minor Go coding issues

internal/attacktechniques/azure/execution/vm-custom-script-extension/main.go

@@ -3,6 +3,8 @@ package azure
 import (
 	"context"
 	_ "embed"
+	"errors"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"log"
 	"time"
 
@@ -38,6 +40,7 @@ Detonation:
 `,
 		Detection:                  "Identify `Microsoft.Compute/virtualMachines/extensions/write` events in Azure Activity logs",
 		Platform:                   stratus.Azure,
+		IsSlow:                     true,
 		IsIdempotent:               false,
 		MitreAttackTactics:         []mitreattack.Tactic{mitreattack.Execution},
 		PrerequisitesTerraformCode: tf,
@@ -46,6 +49,8 @@ Detonation:
 	})
 }
 
+const ExtensionName = "CustomScriptExtension-Stratus-Example"
+
 func detonate(params map[string]string) error {
 	vmObjectId := params["vm_instance_object_id"]
 	vmName := params["vm_name"]
@@ -58,40 +63,53 @@ func detonate(params map[string]string) error {
 
 	client, err := armcompute.NewVirtualMachineExtensionsClient(subscriptionID, cred, clientOptions)
 	if err != nil {
-		log.Fatalf("failed to create client: %v", err)
+		return errors.New("failed to create client: " + err.Error())
 	}
 
 	log.Println("Configuring Custom Script Extension for VM instance " + vmObjectId)
 
+	vmExtension := armcompute.VirtualMachineExtension{
+		Location: to.Ptr("West US"),
+		Properties: &armcompute.VirtualMachineExtensionProperties{
+			Type:                    to.Ptr("CustomScriptExtension"),
+			AutoUpgradeMinorVersion: to.Ptr(true),
+			EnableAutomaticUpgrade:  to.Ptr(false),
+			ProtectedSettings: map[string]interface{}{
+				"commandToExecute": "powershell.exe Get-Service", // the powershell to run with the custom script extension
+			},
+			Publisher:          to.Ptr("Microsoft.Compute"),
+			Settings:           map[string]interface{}{},
+			SuppressFailures:   to.Ptr(true),
+			TypeHandlerVersion: to.Ptr("1.9"),
+		},
+	}
+
 	poller, err := client.BeginCreateOrUpdate(ctx,
 		resourceGroup,
 		vmName,
-		"CustomScriptExtension-Stratus-Example",
-		armcompute.VirtualMachineExtension{
-			Location: to.Ptr("West US"),
-			Properties: &armcompute.VirtualMachineExtensionProperties{
-				Type:                    to.Ptr("CustomScriptExtension"),
-				AutoUpgradeMinorVersion: to.Ptr(true),
-				EnableAutomaticUpgrade:  to.Ptr(false),
-				ProtectedSettings: map[string]interface{}{
-					"commandToExecute": "powershell.exe Get-Service", // the powershell to run with the custom script extension
-				},
-				Publisher:          to.Ptr("Microsoft.Compute"),
-				Settings:           map[string]interface{}{},
-				SuppressFailures:   to.Ptr(true),
-				TypeHandlerVersion: to.Ptr("1.9"),
-			},
-		},
-		&armcompute.VirtualMachineExtensionsClientBeginCreateOrUpdateOptions{ResumeToken: ""})
+		ExtensionName,
+		vmExtension,
+		nil)
+
 	if err != nil {
-		log.Fatalf("failed to finish the request: %v", err)
+		return errors.New("unable to create virtual machine extension: " + err.Error())
 	}
-	res, err := poller.PollUntilDone(ctx, 30*time.Second)
+
+	log.Println("Waiting for Custom Script Extension to be installed on the VM")
+	ctxWithTimeout, done := context.WithTimeout(context.Background(), 60*3*time.Second)
+	defer done()
+	_, err = poller.PollUntilDone(ctxWithTimeout, &runtime.PollUntilDoneOptions{Frequency: 2 * time.Second})
 	if err != nil {
-		log.Fatalf("failed to pull the result: %v", err)
+		return errors.New("unable to retrieve the output of the command ran on the virtual machine: " + err.Error())
 	}
 
-	_ = res
+	/*ctxWithTimeout, done = context.WithTimeout(context.Background(), 60*3*time.Second)
+	defer done()
+	client2, _ := armcompute.NewVirtualMachinesClient(subscriptionID, cred, clientOptions)
+	const tpe = armcompute.InstanceViewTypes()
+	result, _ := client2.Get(ctxWithTimeout, resourceGroup, vmName, &armcompute.VirtualMachinesClientGetOptions{Expand: &tpe})
+	fmt.Println(result.VirtualMachine.Resources[0].Properties.InstanceView.Substatuses[0].Message)
+	return nil*/
 
 	return nil
 }
@@ -116,16 +134,19 @@ func revert(params map[string]string) error {
 	poller, err := client.BeginDelete(ctx,
 		resourceGroup,
 		vmName,
-		"CustomScriptExtension-Stratus-Example",
+		ExtensionName,
 		&armcompute.VirtualMachineExtensionsClientBeginDeleteOptions{ResumeToken: ""})
 
 	if err != nil {
-		log.Fatalf("failed to finish the request: %v", err)
+		return errors.New("unable to remove custom script extension: " + err.Error())
 	}
 
-	_, err = poller.PollUntilDone(ctx, 30*time.Second)
+	ctxWithTimeout, done := context.WithTimeout(context.Background(), 60*3*time.Second)
+	defer done()
+
+	_, err = poller.PollUntilDone(ctxWithTimeout, &runtime.PollUntilDoneOptions{Frequency: 2 * time.Second})
 	if err != nil {
-		log.Fatalf("failed to pull the result: %v", err)
+		return errors.New("unable to remove custom script extension: " + err.Error())
 	}
 
 	return nil