@@ -24,7 +24,11 @@ import (
2424 "net/http"
2525 "net/url"
2626 "sync"
27+ "time"
2728
29+ "github.com/Azure/azure-sdk-for-go/sdk/azcore"
30+ "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
31+ "github.com/Azure/azure-sdk-for-go/sdk/azidentity"
2832 "github.com/bloodhoundad/azurehound/v2/client/config"
2933 "github.com/bloodhoundad/azurehound/v2/constants"
3034)
@@ -52,6 +56,14 @@ type ManagedIdentityAuthStrategy struct {
5256 token Token
5357}
5458
59+ // ManagedIdentitySDKAuthStrategy is an authentication strategy that uses Azure Managed Identity SDK
60+ type ManagedIdentitySDKAuthStrategy struct {
61+ credential * azidentity.ManagedIdentityCredential
62+ resource url.URL
63+ token azcore.AccessToken
64+ mutex sync.Mutex
65+ }
66+
5567// GenericAuthStrategy is an authentication strategy that uses a bunch of pre-existing authentication methods (TODO: Break this up)
5668type GenericAuthStrategy struct {
5769 config config.Config
@@ -83,6 +95,32 @@ func NewManagedIdentityAuthenticator(config config.Config, auth *url.URL, api *u
8395 }
8496}
8597
98+ func GetManagedIdentityCredential (config config.Config ) (* azidentity.ManagedIdentityCredential , error ) {
99+ options := & azidentity.ManagedIdentityCredentialOptions {
100+ ID : azidentity .ClientID (config .ManagedIdentityClientId ),
101+ }
102+ cred , err := azidentity .NewManagedIdentityCredential (options )
103+ if err != nil {
104+ optionsFallback := & azidentity.ManagedIdentityCredentialOptions {}
105+ cred , err = azidentity .NewManagedIdentityCredential (optionsFallback )
106+ if err != nil {
107+ return nil , fmt .Errorf ("failed to authenticate with managed identity: %w" , err )
108+ }
109+ }
110+ return cred , nil
111+ }
112+
113+ // NewManagedIdentitySDKAuthenticator creates a new Authenticator using the ManagedIdentitySDKAuthStrategy
114+ func NewManagedIdentitySDKAuthenticator (config config.Config , resource * url.URL , cred * azidentity.ManagedIdentityCredential ) * Authenticator {
115+ return & Authenticator {
116+ auth : & ManagedIdentitySDKAuthStrategy {
117+ credential : cred ,
118+ resource : * resource ,
119+ },
120+ mutex : sync.RWMutex {},
121+ }
122+ }
123+
86124// NewGenericAuthenticator creates a new Authenticator using the GenericAuthStrategy (The collection of pre-existing authentication methods)
87125func NewGenericAuthenticator (config config.Config , auth * url.URL , api * url.URL ) * Authenticator {
88126 return & Authenticator {
@@ -123,9 +161,14 @@ func (s *Authenticator) refreshIfExpired(r *restClient) error {
123161 return nil
124162 }
125163 // Authenticate
126- if authRequest , err := s .auth .createAuthRequest (); err != nil {
164+ authRequest , err := s .auth .createAuthRequest ()
165+ if err != nil {
127166 return err
128- } else if authResponse , err := r .send (authRequest ); err != nil {
167+ }
168+ if authRequest == nil {
169+ return nil
170+ }
171+ if authResponse , err := r .send (authRequest ); err != nil {
129172 return err
130173 } else {
131174 defer authResponse .Body .Close ()
@@ -249,3 +292,43 @@ func (s *GenericAuthStrategy) addAuthenticationToRequest(req *http.Request) (*ht
249292 }
250293 return req , nil
251294}
295+
296+ // Adds the access token to the outgoing HTTP request
297+ func (s * ManagedIdentitySDKAuthStrategy ) addAuthenticationToRequest (req * http.Request ) (* http.Request , error ) {
298+ token := s .token
299+ s .mutex .Lock ()
300+ defer s .mutex .Unlock ()
301+ if token .Token == "" || token .ExpiresOn .Before (time .Now ().Add (2 * time .Minute )) {
302+ if err := s .refreshToken (req .Context ()); err != nil {
303+ return nil , err
304+ }
305+ token = s .token
306+ }
307+ req .Header .Set ("Authorization" , "Bearer " + token .Token )
308+ return req , nil
309+ }
310+
311+ func (s * ManagedIdentitySDKAuthStrategy ) refreshToken (ctx context.Context ) error {
312+ token , err := s .credential .GetToken (ctx , policy.TokenRequestOptions {
313+ Scopes : []string {s .resource .String () + "/.default" },
314+ })
315+ if err != nil {
316+ return err
317+ }
318+ s .token = token
319+ return nil
320+ }
321+
322+ func (s * ManagedIdentitySDKAuthStrategy ) createAuthRequest () (* http.Request , error ) {
323+ // Not used in SDK-based flow, return nil
324+ return nil , nil
325+ }
326+
327+ func (s * ManagedIdentitySDKAuthStrategy ) decodeAuthResponse (resp * http.Response ) error {
328+ // Not used in SDK-based flow, do nothing
329+ return nil
330+ }
331+
332+ func (s * ManagedIdentitySDKAuthStrategy ) isExpired () bool {
333+ return s .token .Token == "" || s .token .ExpiresOn .Before (time .Now ().Add (2 * time .Minute ))
334+ }
0 commit comments