Add ability for client to use delegated token authentication (DataDog#3133)

* Add ability for client to use delegated token authentication in place of api and app key authentication

* Cleaned up code quality violation

* Trim the api prefix if it is present in the name serverVars

* Added tests and general code cleanup

* Updated how we pass the token to the subsequent requests

* pre-commit fixes

* Update code generation for client

* pre-commit fixes

* Fixed non-standard indentation

* Refactored to structure to put delegated auth into the datadog package. Added templates for aws.go and delegated_auth.go

* Fixed test

* Refactored how we get the server URL for delegated auth to use the standard ServerConfigurations

* pre-commit fixes

* Cleaned up generator and test

* AUTHN-4821 - Refactored delegated auth to use not use the context as much and not to hijack the SetAuthKeys

* AUTHN-4821 - Cleaned up template code

* pre-commit fixes

* AUTHN-4821 - Updated AWS variable naming

* AUTHN-4821 - Fixed issues with context updating

* AUTHN-4821 - Fixed AWS example

* AUTHN-4821 - Code cleanup, comments and moving a call

* pre-commit fixes

---------

Co-authored-by: ci.datadog-api-spec <packages@datadoghq.com>

Author: wynbennett

.generator/conftest.py

@@ -1,30 +1,26 @@
 # coding=utf-8
 """Define basic fixtures."""
 
+import hashlib
 import json
 import os
 import pathlib
+import pytest
 import re
 import warnings
 import zlib
 from collections import defaultdict
-
-import pytest
 from dateutil.relativedelta import relativedelta
-from jinja2 import Environment, FileSystemLoader, Template
-from pytest_bdd import given, parsers, then, when
-import hashlib
-
 from generator import openapi
+from generator.formatter import format_parameters, format_data_with_schema, go_name
 from generator.utils import (
     camel_case,
     given_variables,
     snake_case,
     untitle_case,
 )
-
-from generator.formatter import format_parameters, format_data_with_schema, go_name
-
+from jinja2 import Environment, FileSystemLoader, Template
+from pytest_bdd import given, parsers, then, when
 
 MODIFIED_FEATURES = {pathlib.Path(p).resolve() for p in os.getenv("BDD_MODIFIED_FEATURES", "").split(" ") if p}
 
@@ -74,6 +70,9 @@ def lookup(value, path):
 JINJA_ENV.globals["given_variables"] = given_variables
 
 GO_EXAMPLE_J2 = JINJA_ENV.get_template("example.j2")
+DATADOG_EXAMPLES_J2 = {
+    "aws.go": JINJA_ENV.get_template("example_aws.j2")
+}
 
 
 def pytest_bdd_after_scenario(request, feature, scenario):
@@ -113,6 +112,19 @@ def pytest_bdd_after_scenario(request, feature, scenario):
     with output.open("w") as f:
         f.write(data)
 
+    for file_name, template in DATADOG_EXAMPLES_J2.items():
+        output = ROOT_PATH / "examples" / "datadog" / file_name
+        output.parent.mkdir(parents=True, exist_ok=True)
+
+        data = template.render(
+            context=context,
+            version=version,
+            scenario=scenario,
+            operation_spec=operation_spec.spec,
+        )
+        with output.open("w") as f:
+            f.write(data)
+
 
 def pytest_bdd_apply_tag(tag, function):
     """Register tags as custom markers and skip test for '@skip' ones."""

.generator/src/generator/cli.py

@@ -1,10 +1,9 @@
-import pathlib
-
 import click
+import pathlib
 from jinja2 import Environment, FileSystemLoader
 
-from . import openapi
 from . import formatter
+from . import openapi
 from . import utils
 
 MODULE = "github.com/DataDog/datadog-api-client-go/v2"
@@ -67,8 +66,10 @@ def cli(specs, output):
     doc_j2 = env.get_template("doc.j2")
 
     extra_files = {
+        "aws.go": env.get_template("aws.j2"),
         "client.go": env.get_template("client.j2"),
         "configuration.go": env.get_template("configuration.j2"),
+        "delegated_auth.go": env.get_template("delegated_auth.j2"),
         "utils.go": env.get_template("utils.j2"),
         "zstd.go": env.get_template("zstd.j2"),
         "no_zstd.go": env.get_template("no_zstd.j2"),

.generator/src/generator/templates/api.j2

@@ -241,7 +241,23 @@ localVarQueryParams.Add("{{ parameter.name }}", {{ common_package_name }}.Parame
 	{%- endif %}
 
 	{%- set authMethods = operation.security if "security" in operation else openapi.security %}
+	{%- set appKeyNs = namespace(hasAppKeyAuth=false) %}
+	{%- for authMethod in authMethods %}
+		{%- for name in authMethod %}
+			{%- if name == "appKeyAuth" %}
+				{%- set appKeyNs.hasAppKeyAuth = true %}
+			{%- endif %}
+		{%- endfor %}
+	{%- endfor %}
 	{%- if authMethods %}
+	{%- if appKeyNs.hasAppKeyAuth %}
+		if a.Client.Cfg.DelegatedTokenConfig != nil {
+			err = {{ common_package_name }}.UseDelegatedTokenAuth(ctx, &localVarHeaderParams, a.Client.Cfg.DelegatedTokenConfig)
+			if err != nil {
+				return {% if returnType %}localVarReturnValue, {% endif %}nil, err
+			}
+		} else {
+	{%- endif %}
         {{ common_package_name }}.SetAuthKeys(
 		ctx,
 		&localVarHeaderParams,
@@ -254,6 +270,7 @@ localVarQueryParams.Add("{{ parameter.name }}", {{ common_package_name }}.Parame
 	{%- endfor %}
 	{%- endfor %}
 	)
+	{% if appKeyNs.hasAppKeyAuth %} } {% endif %}
 	{%- endif %}
 	req, err := a.Client.PrepareRequest(ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, {% if formParameter %}&formFile{% else %}nil{% endif %})
 	if err != nil {

.generator/src/generator/templates/aws.j2

@@ -0,0 +1,252 @@
+{% include "partial_header.j2" %}
+package {{ common_package_name }}
+
+import (
+	"context"
+	"crypto/hmac"
+	"crypto/sha256"
+	"encoding/base64"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"os"
+	"sort"
+	"strings"
+	"time"
+)
+
+type Credentials struct {
+	AccessKeyID     string
+	SecretAccessKey string
+	SessionToken    string
+}
+
+// SigningData is the data structure that represents the Data used to generate and AWS Proof
+type SigningData struct {
+	HeadersEncoded string `json:"iam_headers_encoded"`
+	BodyEncoded    string `json:"iam_body_encoded"`
+	URLEncoded     string `json:"iam_url_encoded"`
+	Method         string `json:"iam_method"`
+}
+
+const (
+	// Common Headers
+	// orgIdHeader is the header we use to specify the name of the org we request a token for
+	orgIdHeader     = "x-ddog-org-id"
+	hostHeader      = "host"
+	applicationForm = "application/x-www-form-urlencoded; charset=utf-8"
+
+	// AWS specific constants
+	AWSAccessKeyIdName     = "AWS_ACCESS_KEY_ID"
+	AWSSecretAccessKeyName = "AWS_SECRET_ACCESS_KEY"
+	AWSSessionTokenName    = "AWS_SESSION_TOKEN"
+
+	amzDateHeader         = "X-Amz-Date"
+	amzTokenHeader        = "X-Amz-Security-Token"
+	amzDateFormat         = "20060102"
+	amzDateTimeFormat     = "20060102T150405Z"
+	defaultRegion         = "us-east-1"
+	defaultStsHost        = "sts.amazonaws.com"
+	regionalStsHost       = "sts.%s.amazonaws.com"
+	service               = "sts"
+	algorithm             = "AWS4-HMAC-SHA256"
+	aws4Request           = "aws4_request"
+	getCallerIdentityBody = "Action=GetCallerIdentity&Version=2011-06-15"
+)
+
+const ProviderAWS = "aws"
+
+type AWSAuth struct {
+	AwsRegion string
+}
+
+func (a *AWSAuth) Authenticate(ctx context.Context, config *DelegatedTokenConfig) (*DelegatedTokenCredentials, error) {
+	// Get local AWS Credentials
+	creds := a.GetCredentials(ctx)
+
+	if config == nil || config.OrgUUID == "" {
+		return nil, fmt.Errorf("missing org UUID in config")
+	}
+
+	// Use the credentials to generate the signing data
+	data, err := a.GenerateAwsAuthData(config.OrgUUID, creds)
+	if err != nil {
+		return nil, err
+	}
+
+	// Generate the auth string passed to the token endpoint
+	authString := data.BodyEncoded + "|" + data.HeadersEncoded + "|" + data.Method + "|" + data.URLEncoded
+
+	authResponse, err := GetDelegatedToken(ctx, config.OrgUUID, authString)
+	return authResponse, err
+}
+
+func (a *AWSAuth) GetCredentials(ctx context.Context) *Credentials {
+	keys := ctx.Value(ContextAWSVariables)
+	if keys != nil {
+		keysMap := keys.(map[string]string)
+		creds := Credentials{}
+		if accessKey, ok := keysMap[AWSAccessKeyIdName]; ok {
+			creds.AccessKeyID = accessKey
+		}
+		if secretKey, ok := keysMap[AWSSecretAccessKeyName]; ok {
+			creds.SecretAccessKey = secretKey
+		}
+		if sessionToken, ok := keysMap[AWSSessionTokenName]; ok {
+			creds.SessionToken = sessionToken
+		}
+		return &creds
+	} else {
+		accessKey := os.Getenv(AWSAccessKeyIdName)
+		secretKey := os.Getenv(AWSSecretAccessKeyName)
+		sessionToken := os.Getenv(AWSSessionTokenName)
+		return &Credentials{
+			AccessKeyID:     accessKey,
+			SecretAccessKey: secretKey,
+			SessionToken:    sessionToken,
+		}
+	}
+}
+
+func (a *AWSAuth) getConnectionParameters() (string, string, string) {
+	region := a.AwsRegion
+	var host string
+	// Default to the default global STS Host (see here: https://docs.aws.amazon.com/general/latest/gr/sts.html)
+	if region == "" {
+		region = defaultRegion
+		host = defaultStsHost
+	} else {
+		// If the region is not empty, use the regional STS host
+		host = fmt.Sprintf(regionalStsHost, region)
+	}
+	stsFullURL := fmt.Sprintf("https://%s", host)
+	return stsFullURL, region, host
+}
+
+func (a *AWSAuth) GenerateAwsAuthData(orgUUID string, creds *Credentials) (*SigningData, error) {
+	if orgUUID == "" {
+		return nil, fmt.Errorf("missing org UUID")
+	}
+	if creds == nil || (creds.AccessKeyID == "" && creds.SecretAccessKey == "") || creds.SessionToken == "" {
+		return nil, fmt.Errorf("missing AWS credentials")
+	}
+	stsFullURL, region, host := a.getConnectionParameters()
+
+	now := time.Now().UTC()
+
+	requestBody := getCallerIdentityBody
+	h := sha256.Sum256([]byte(requestBody))
+	payloadHash := hex.EncodeToString(h[:])
+
+	// Create the headers that factor into the signing algorithm
+	headerMap := map[string][]string{
+		contextLengthHeader: {
+			fmt.Sprintf("%d", len(requestBody)),
+		},
+		contentTypeHeader: {
+			applicationForm,
+		},
+		amzDateHeader: {
+			now.Format(amzDateTimeFormat),
+		},
+		orgIdHeader: {
+			orgUUID,
+		},
+		amzTokenHeader: {
+			creds.SessionToken,
+		},
+		hostHeader: {
+			host,
+		},
+	}
+
+	headerArr := make([]string, len(headerMap), len(headerMap))
+	signedHeadersArr := make([]string, len(headerMap), len(headerMap))
+	headerIdx := 0
+	for k, v := range headerMap {
+		loweredHeaderName := strings.ToLower(k)
+		headerArr[headerIdx] = fmt.Sprintf("%s:%s", loweredHeaderName, strings.Join(v, ","))
+		signedHeadersArr[headerIdx] = loweredHeaderName
+		headerIdx++
+	}
+	sort.Strings(headerArr)
+	sort.Strings(signedHeadersArr)
+	signedHeaders := strings.Join(signedHeadersArr, ";")
+
+	canonicalRequest := strings.Join([]string{
+		http.MethodPost,
+		"/",
+		"", // No query string
+		strings.Join(headerArr, "\n") + "\n",
+		signedHeaders,
+		payloadHash,
+	}, "\n")
+
+	// Create the string to sign
+	hashCanonicalRequest := sha256.Sum256([]byte(canonicalRequest))
+	credentialScope := strings.Join([]string{
+		now.Format(amzDateFormat),
+		region,
+		service,
+		aws4Request,
+	}, "/")
+	stringToSign := a.makeSignature(
+		now,
+		credentialScope,
+		hex.EncodeToString(hashCanonicalRequest[:]),
+		region,
+		service,
+		creds.SecretAccessKey,
+		algorithm,
+	)
+
+	// Create the authorization header
+	credential := strings.Join([]string{
+		creds.AccessKeyID,
+		credentialScope,
+	}, "/")
+	authHeader := fmt.Sprintf("%s Credential=%s, SignedHeaders=%s, Signature=%s",
+		algorithm, credential, signedHeaders, stringToSign)
+
+	headerMap["Authorization"] = []string{authHeader}
+	headerMap["User-Agent"] = []string{GetUserAgent()}
+	headersJSON, err := json.Marshal(headerMap)
+	if err != nil {
+		return nil, err
+	}
+
+	return &SigningData{
+		HeadersEncoded: base64.StdEncoding.EncodeToString(headersJSON),
+		BodyEncoded:    base64.StdEncoding.EncodeToString([]byte(requestBody)),
+		Method:         http.MethodPost,
+		URLEncoded:     base64.StdEncoding.EncodeToString([]byte(stsFullURL)),
+	}, nil
+}
+
+func (a *AWSAuth) makeSignature(t time.Time, credentialScope, payloadHash, region, service, secretAccessKey, algorithm string) string {
+	// Create the string to sign
+	stringToSign := strings.Join([]string{
+		algorithm,
+		t.Format(amzDateTimeFormat),
+		credentialScope,
+		payloadHash,
+	}, "\n")
+
+	// Create the signing key
+	kDate := hmac256(t.Format(amzDateFormat), []byte("AWS4"+secretAccessKey))
+	kRegion := hmac256(region, kDate)
+	kService := hmac256(service, kRegion)
+	kSigning := hmac256(aws4Request, kService)
+
+	// Sign the string
+	signature := hex.EncodeToString(hmac256(stringToSign, kSigning))
+
+	return signature
+}
+
+func hmac256(data string, key []byte) []byte {
+	h := hmac.New(sha256.New, key)
+	h.Write([]byte(data))
+	return h.Sum(nil)
+}