Add more parameters to get-access-token (#643)

* Add more parameters to get-access-token

* Fix comments

* Rename GetTokenFromCLIWithParams & Expose Azure CLI directly

* Change func name

* Fix parameter validation

* Change comment

* Fix comment

Co-authored-by: Joel Hendrix <jhendrix@microsoft.com>

autorest/azure/cli/token.go

@@ -124,38 +124,46 @@ func LoadTokens(path string) ([]Token, error) {
 
 // GetTokenFromCLI gets a token using Azure CLI 2.0 for local development scenarios.
 func GetTokenFromCLI(resource string) (*Token, error) {
-	// This is the path that a developer can set to tell this class what the install path for Azure CLI is.
-	const azureCLIPath = "AzureCLIPath"
+	return GetTokenFromCLIWithParams(GetAccessTokenParams{Resource: resource})
+}
 
-	// The default install paths are used to find Azure CLI. This is for security, so that any path in the calling program's Path environment is not used to execute Azure CLI.
-	azureCLIDefaultPathWindows := fmt.Sprintf("%s\\Microsoft SDKs\\Azure\\CLI2\\wbin; %s\\Microsoft SDKs\\Azure\\CLI2\\wbin", os.Getenv("ProgramFiles(x86)"), os.Getenv("ProgramFiles"))
+// GetAccessTokenParams is the parameter struct of GetTokenFromCLIWithParams
+type GetAccessTokenParams struct {
+	Resource     string
+	ResourceType string
+	Subscription string
+	Tenant       string
+}
 
-	// Default path for non-Windows.
-	const azureCLIDefaultPath = "/bin:/sbin:/usr/bin:/usr/local/bin"
+// GetTokenFromCLIWithParams gets a token using Azure CLI 2.0 for local development scenarios.
+func GetTokenFromCLIWithParams(params GetAccessTokenParams) (*Token, error) {
+	cliCmd := GetAzureCLICommand()
 
-	// Validate resource, since it gets sent as a command line argument to Azure CLI
-	const invalidResourceErrorTemplate = "Resource %s is not in expected format. Only alphanumeric characters, [dot], [colon], [hyphen], and [forward slash] are allowed."
-	match, err := regexp.MatchString("^[0-9a-zA-Z-.:/]+$", resource)
-	if err != nil {
-		return nil, err
+	cliCmd.Args = append(cliCmd.Args, "account", "get-access-token", "-o", "json")
+	if params.Resource != "" {
+		if err := validateParameter(params.Resource); err != nil {
+			return nil, err
+		}
+		cliCmd.Args = append(cliCmd.Args, "--resource", params.Resource)
 	}
-	if !match {
-		return nil, fmt.Errorf(invalidResourceErrorTemplate, resource)
+	if params.ResourceType != "" {
+		if err := validateParameter(params.ResourceType); err != nil {
+			return nil, err
+		}
+		cliCmd.Args = append(cliCmd.Args, "--resource-type", params.ResourceType)
 	}
-
-	// Execute Azure CLI to get token
-	var cliCmd *exec.Cmd
-	if runtime.GOOS == "windows" {
-		cliCmd = exec.Command(fmt.Sprintf("%s\\system32\\cmd.exe", os.Getenv("windir")))
-		cliCmd.Env = os.Environ()
-		cliCmd.Env = append(cliCmd.Env, fmt.Sprintf("PATH=%s;%s", os.Getenv(azureCLIPath), azureCLIDefaultPathWindows))
-		cliCmd.Args = append(cliCmd.Args, "/c", "az")
-	} else {
-		cliCmd = exec.Command("az")
-		cliCmd.Env = os.Environ()
-		cliCmd.Env = append(cliCmd.Env, fmt.Sprintf("PATH=%s:%s", os.Getenv(azureCLIPath), azureCLIDefaultPath))
+	if params.Subscription != "" {
+		if err := validateParameter(params.Subscription); err != nil {
+			return nil, err
+		}
+		cliCmd.Args = append(cliCmd.Args, "--subscription", params.Subscription)
+	}
+	if params.Tenant != "" {
+		if err := validateParameter(params.Tenant); err != nil {
+			return nil, err
+		}
+		cliCmd.Args = append(cliCmd.Args, "--tenant", params.Tenant)
 	}
-	cliCmd.Args = append(cliCmd.Args, "account", "get-access-token", "-o", "json", "--resource", resource)
 
 	var stderr bytes.Buffer
 	cliCmd.Stderr = &stderr
@@ -173,3 +181,43 @@ func GetTokenFromCLI(resource string) (*Token, error) {
 
 	return &tokenResponse, err
 }
+
+func validateParameter(param string) error {
+	// Validate parameters, since it gets sent as a command line argument to Azure CLI
+	const invalidResourceErrorTemplate = "Parameter %s is not in expected format. Only alphanumeric characters, [dot], [colon], [hyphen], and [forward slash] are allowed."
+	match, err := regexp.MatchString("^[0-9a-zA-Z-.:/]+$", param)
+	if err != nil {
+		return err
+	}
+	if !match {
+		return fmt.Errorf(invalidResourceErrorTemplate, param)
+	}
+	return nil
+}
+
+// GetAzureCLICommand can be used to run arbitrary Azure CLI command
+func GetAzureCLICommand() *exec.Cmd {
+	// This is the path that a developer can set to tell this class what the install path for Azure CLI is.
+	const azureCLIPath = "AzureCLIPath"
+
+	// The default install paths are used to find Azure CLI. This is for security, so that any path in the calling program's Path environment is not used to execute Azure CLI.
+	azureCLIDefaultPathWindows := fmt.Sprintf("%s\\Microsoft SDKs\\Azure\\CLI2\\wbin; %s\\Microsoft SDKs\\Azure\\CLI2\\wbin", os.Getenv("ProgramFiles(x86)"), os.Getenv("ProgramFiles"))
+
+	// Default path for non-Windows.
+	const azureCLIDefaultPath = "/bin:/sbin:/usr/bin:/usr/local/bin"
+
+	// Execute Azure CLI to get token
+	var cliCmd *exec.Cmd
+	if runtime.GOOS == "windows" {
+		cliCmd = exec.Command(fmt.Sprintf("%s\\system32\\cmd.exe", os.Getenv("windir")))
+		cliCmd.Env = os.Environ()
+		cliCmd.Env = append(cliCmd.Env, fmt.Sprintf("PATH=%s;%s", os.Getenv(azureCLIPath), azureCLIDefaultPathWindows))
+		cliCmd.Args = append(cliCmd.Args, "/c", "az")
+	} else {
+		cliCmd = exec.Command("az")
+		cliCmd.Env = os.Environ()
+		cliCmd.Env = append(cliCmd.Env, fmt.Sprintf("PATH=%s:%s", os.Getenv(azureCLIPath), azureCLIDefaultPath))
+	}
+
+	return cliCmd
+}