-
Notifications
You must be signed in to change notification settings - Fork 379
/
azure_cli_auth.go
115 lines (105 loc) · 3.31 KB
/
azure_cli_auth.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
package common
import (
"context"
"encoding/json"
"fmt"
"log"
"net/http"
"os/exec"
"strings"
"sync"
"time"
"github.com/Azure/go-autorest/autorest"
"github.com/Azure/go-autorest/autorest/adal"
"github.com/Azure/go-autorest/autorest/azure/cli"
)
type refreshableCliToken struct {
resource string
token *adal.Token
lock *sync.RWMutex
refreshMinutes int
}
// OAuthToken implements adal.OAuthTokenProvider
func (rct *refreshableCliToken) OAuthToken() string {
if rct.token == nil {
return ""
}
return rct.token.OAuthToken()
}
// EnsureFreshWithContext implements adal.RefresherWithContext
func (rct *refreshableCliToken) EnsureFreshWithContext(ctx context.Context) error {
refreshInterval := time.Duration(rct.refreshMinutes) * time.Minute
if rct.token != nil && !rct.token.WillExpireIn(refreshInterval) {
return nil
}
rct.lock.Lock()
defer rct.lock.Unlock()
if rct.token != nil && !rct.token.WillExpireIn(refreshInterval) {
return nil
}
return rct.refreshInternal(rct.resource)
}
// RefreshWithContext implements adal.RefresherWithContext
func (rct *refreshableCliToken) RefreshWithContext(ctx context.Context) error {
rct.lock.Lock()
defer rct.lock.Unlock()
return rct.refreshInternal(rct.resource)
}
// RefreshExchangeWithContext implements adal.RefresherWithContext
func (rct *refreshableCliToken) RefreshExchangeWithContext(ctx context.Context, resource string) error {
rct.lock.Lock()
defer rct.lock.Unlock()
return rct.refreshInternal(rct.resource)
}
func (rct *refreshableCliToken) refreshInternal(resource string) error {
out, err := exec.Command("az", "account", "get-access-token", "--resource", resource, "--output", "json").Output()
if ee, ok := err.(*exec.ExitError); ok {
return fmt.Errorf("cannot get access token: %s", string(ee.Stderr))
}
if err != nil {
return fmt.Errorf("cannot get access token: %v", err)
}
var cliToken cli.Token
err = json.Unmarshal(out, &cliToken)
if err != nil {
return fmt.Errorf("cannot unmarshal CLI result: %w", err)
}
token, err := cliToken.ToADALToken()
if err != nil {
return fmt.Errorf("cannot convert to ADAL token: %w", err)
}
log.Printf("[INFO] Refreshed OAuth token for %s from Azure CLI, which expires on %s", resource, cliToken.ExpiresOn)
rct.token = &token
return nil
}
func (aa *DatabricksClient) cliAuthorizer(resource string) (autorest.Authorizer, error) {
rct := refreshableCliToken{
lock: &sync.RWMutex{},
resource: resource,
refreshMinutes: 6,
}
err := rct.refreshInternal(resource)
if err != nil {
return nil, fmt.Errorf("cannot refresh: %w", err)
}
return autorest.NewBearerAuthorizer(&rct), nil
}
func (aa *DatabricksClient) configureWithAzureCLI(ctx context.Context) (func(*http.Request) error, error) {
if !aa.IsAzure() {
return nil, nil
}
if aa.IsAzureClientSecretSet() {
return nil, nil
}
// verify that Azure CLI is authenticated
_, err := cli.GetTokenFromCLI(armDatabricksResourceID)
if err != nil {
if strings.Contains(err.Error(), "executable file not found") {
return nil, fmt.Errorf("most likely Azure CLI is not installed. " +
"See https://docs.microsoft.com/en-us/cli/azure/?view=azure-cli-latest for details")
}
return nil, err
}
log.Printf("[INFO] Using Azure CLI authentication with AAD tokens")
return aa.simpleAADRequestVisitor(ctx, aa.cliAuthorizer)
}