Merge pull request #103 from scaryrawr/fix/az-shim-access-token

fix: align az get-access-token shim behavior

Author: markphip

src/artifacts-helper/devcontainer-feature.json

@@ -1,7 +1,7 @@
 {
     "name": "Azure Artifacts Credential Helper",
     "id": "artifacts-helper",
-    "version": "3.0.4",
+    "version": "3.0.5",
     "description": "Configures Codespace to authenticate with Azure Artifact feeds",
     "options": {
         "nugetURIPrefixes": {
@@ -85,4 +85,4 @@
             ]
         }
     }
-}
+}

src/artifacts-helper/scripts/az

@@ -28,18 +28,25 @@ if [[ "$1" == "account" && "$2" == "get-access-token" ]]; then
     resource=""
     scope=""
     resource_type=""
+    query=""
+    output_format=""
     prev=""
 
     for arg in "${@:3}"; do
         case "$arg" in
             --resource=*) resource="${arg#--resource=}" ;;
             --scope=*) scope="${arg#--scope=}" ;;
             --resource-type=*) resource_type="${arg#--resource-type=}" ;;
+            --query=*) query="${arg#--query=}" ;;
+            --output=*) output_format="${arg#--output=}" ;;
+            -o=*) output_format="${arg#-o=}" ;;
             *)
                 case "$prev" in
                     --resource) resource="$arg" ;;
                     --scope) scope="$arg" ;;
                     --resource-type) resource_type="$arg" ;;
+                    --query|-q) query="$arg" ;;
+                    --output|-o) output_format="$arg" ;;
                 esac
                 ;;
         esac
@@ -51,6 +58,11 @@ if [[ "$1" == "account" && "$2" == "get-access-token" ]]; then
         resource="${RESOURCE_TYPE_MAP[$resource_type]}"
     fi
 
+    # Default to Azure DevOps resource if no resource or scope specified
+    if [[ -z "$resource" && -z "$scope" ]]; then
+        resource="499b84ac-1321-427f-aa17-267ca6975798"
+    fi
+
     # Determine the scope to request
     request_scope=""
     if [[ -n "$scope" ]]; then
@@ -67,6 +79,18 @@ if [[ "$1" == "account" && "$2" == "get-access-token" ]]; then
     if [[ -n "$request_scope" && -f "${HOME}/azure-auth-helper" ]]; then
         token=$("${HOME}/azure-auth-helper" get-access-token "$request_scope" 2>/dev/null)
         if [[ $? -eq 0 && -n "$token" ]]; then
+            # Handle --query accessToken for direct token extraction
+            if [[ "$query" == "accessToken" ]]; then
+                if [[ "$output_format" == "tsv" ]]; then
+                    echo "$token"
+                else
+                    escaped_token="${token//\\/\\\\}"
+                    escaped_token="${escaped_token//\"/\\\"}"
+                    echo "\"$escaped_token\""
+                fi
+                exit 0
+            fi
+
             # Escape token for safe JSON embedding (handle backslashes and quotes)
             escaped_token="${token//\\/\\\\}"
             escaped_token="${escaped_token//\"/\\\"}"

test/artifacts-helper/test_az_shim.sh

@@ -63,6 +63,39 @@ HELPER
     echo "SUCCESS" || echo "FAILED"
 ' | grep -q "SUCCESS"
 
+# Test that accessToken queries with TSV output return only the token
+check "az shim honors query accessToken with tsv output" bash -c '
+    export HOME='"$TEST_HOME"'
+    cat > "${HOME}/azure-auth-helper" << '\''HELPER'\''
+#!/bin/bash
+echo "test-token-12345"
+HELPER
+    chmod +x "${HOME}/azure-auth-helper"
+
+    output=$(/usr/local/share/codespace-shims/az account get-access-token --resource https://management.azure.com -q accessToken -o tsv 2>&1)
+    [ "$output" = "test-token-12345" ] && echo "SUCCESS" || echo "FAILED: $output"
+' | grep -q "SUCCESS"
+
+# Test that the shim defaults to the Azure DevOps resource when none is specified
+check "az shim defaults to Azure DevOps resource" bash -c '
+    export HOME='"$TEST_HOME"'
+    cat > "${HOME}/azure-auth-helper" << '\''HELPER'\''
+#!/bin/bash
+if [ "$1" = "get-access-token" ]; then
+    echo "$2" > "${HOME}/requested-scope"
+    echo "default-resource-token"
+fi
+HELPER
+    chmod +x "${HOME}/azure-auth-helper"
+
+    output=$(/usr/local/share/codespace-shims/az account get-access-token -q accessToken -o tsv 2>&1)
+    requested_scope=$(cat "${HOME}/requested-scope")
+
+    [ "$output" = "default-resource-token" ] && \
+    [ "$requested_scope" = "499b84ac-1321-427f-aa17-267ca6975798/.default" ] && \
+    echo "SUCCESS" || echo "FAILED: output=$output scope=$requested_scope"
+' | grep -q "SUCCESS"
+
 # Test GitHub Actions bypass (simulate by setting the env var)
 check "az shim bypasses interception in GitHub Actions" bash -c '
     export HOME='"$TEST_HOME"'