From 0aff276f27ff8de86c6687e26516d7a9184ad8a5 Mon Sep 17 00:00:00 2001
From: M Essam Hamed <github@messam.xyz>
Date: Sun, 25 May 2025 17:14:28 +0300
Subject: [PATCH] [management] REST Client Account Impersonation Support

---
 management/client/rest/client.go  | 17 +++++++++++++----
 management/client/rest/options.go | 30 +++++++++++++++++++++++-------
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/management/client/rest/client.go b/management/client/rest/client.go
index 25e8ad0da..9d07ed73c 100644
--- a/management/client/rest/client.go
+++ b/management/client/rest/client.go
@@ -13,9 +13,10 @@ import (
 
 // Client Management service HTTP REST API Client
 type Client struct {
-	managementURL string
-	authHeader    string
-	httpClient    HttpClient
+	managementURL       string
+	authHeader          string
+	impersonatedAccount string
+	httpClient          HTTPClient
 
 	// Accounts NetBird account APIs
 	// see more: https://docs.netbird.io/api/resources/accounts
@@ -86,7 +87,8 @@ func NewWithBearerToken(managementURL, token string) *Client {
 	)
 }
 
-func NewWithOptions(opts ...option) *Client {
+// NewWithOptions initialize new Client instance with options
+func NewWithOptions(opts ...Option) *Client {
 	client := &Client{
 		httpClient: http.DefaultClient,
 	}
@@ -115,6 +117,7 @@ func (c *Client) initialize() {
 	c.Events = &EventsAPI{c}
 }
 
+// NewRequest creates and executes new management API request
 func (c *Client) NewRequest(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) {
 	req, err := http.NewRequestWithContext(ctx, method, c.managementURL+path, body)
 	if err != nil {
@@ -127,6 +130,12 @@ func (c *Client) NewRequest(ctx context.Context, method, path string, body io.Re
 		req.Header.Add("Content-Type", "application/json")
 	}
 
+	if c.impersonatedAccount != "" {
+		query := req.URL.Query()
+		query.Add("account", c.impersonatedAccount)
+		req.URL.RawQuery = query.Encode()
+	}
+
 	resp, err := c.httpClient.Do(req)
 	if err != nil {
 		return nil, err
diff --git a/management/client/rest/options.go b/management/client/rest/options.go
index 5aad7dd7e..24184d359 100644
--- a/management/client/rest/options.go
+++ b/management/client/rest/options.go
@@ -2,34 +2,50 @@ package rest
 
 import "net/http"
 
-type option func(*Client)
+// Option modifier for creation of Client
+type Option func(*Client)
 
-type HttpClient interface {
+// HTTPClient interface for HTTP client
+type HTTPClient interface {
 	Do(req *http.Request) (*http.Response, error)
 }
 
-func WithHttpClient(client HttpClient) option {
+// WithHTTPClient overrides HTTPClient used
+func WithHTTPClient(client HTTPClient) Option {
 	return func(c *Client) {
 		c.httpClient = client
 	}
 }
 
-func WithBearerToken(token string) option {
+// WithBearerToken uses provided bearer token acquired from SSO for authentication
+func WithBearerToken(token string) Option {
 	return WithAuthHeader("Bearer " + token)
 }
 
-func WithPAT(token string) option {
+// WithPAT uses provided Personal Access Token
+// (created from NetBird Management Dashboard) for authentication
+func WithPAT(token string) Option {
 	return WithAuthHeader("Token " + token)
 }
 
-func WithManagementURL(url string) option {
+// WithManagementURL overrides target NetBird Management server
+func WithManagementURL(url string) Option {
 	return func(c *Client) {
 		c.managementURL = url
 	}
 }
 
-func WithAuthHeader(value string) option {
+// WithAuthHeader overrides auth header completely, this should generally not be used
+// and WithBearerToken or WithPAT should be used instead
+func WithAuthHeader(value string) Option {
 	return func(c *Client) {
 		c.authHeader = value
 	}
 }
+
+// WithAccount uses impersonated account for Client
+func WithAccount(value string) Option {
+	return func(c *Client) {
+		c.impersonatedAccount = value
+	}
+}