Risk API Example

package main

import (
	"context"
	"encoding/csv"
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/pkg/errors"
)


type loginResponse struct {
	AccessTokens map[string]map[string]struct {
		AccessToken string `json:"access_token"`
	} `json:"access_tokens"`
}

type riskSummary struct {
	ID        string `json:"id"`
	Signature struct {
		ID string `json:"id"`
	} `json:"signature"`
}

type listRisksResponse struct {
	Data []riskSummary `json:"data"`
}

type riskResource struct {
	Kind      string `json:"kind"`
	Name      string `json:"name"`
	Namespace string `json:"namespace"`
}

type listRiskResourcesResponse struct {
	Data []riskResource `json:"data"`
}

var (
	clusterID  = flag.String("cluster-id", "", "Chkk cluster ID (required)")
	namespace  = flag.String("namespace", "", "Namespace filter (optional)")
	outFile    = flag.String("out", "risks.csv", "CSV output filename")
	apiBase    = "https://api.us-east-1.aws.chkk.io/v1"
	httpClient = &http.Client{Timeout: 15 * time.Second}
)

func main() {
	flag.Parse()
	if *clusterID == "" {
		exitErr(errors.New("flag -cluster-id is required"))
	}

	ctx := context.Background()

	token, err := authenticate(ctx)
	if err != nil {
		exitErr(err)
	}

	risks, err := listRisks(ctx, token, *clusterID)
	if err != nil {
		exitErr(err)
	}

	if err := printSummary(ctx, token, risks); err != nil {
		exitErr(err)
	}

	if err := writeCSV(ctx, token, risks); err != nil {
		exitErr(err)
	}
}

func printSummary(ctx context.Context, token string, risks []riskSummary) error {
	for _, rs := range risks {
		res, err := listRiskResources(ctx, token, rs.ID)
		if err != nil {
			return err
		}

		count := filterByNamespace(res, *namespace)
		if *namespace != "" {
			fmt.Printf("%s, %s: %d affected resources in namespace %s\n", *clusterID, rs.Signature.ID, count, *namespace)
		} else {
			fmt.Printf("%s, %s: %d affected resources\n", *clusterID, rs.Signature.ID, count)
		}
	}
	return nil
}

func writeCSV(ctx context.Context, token string, risks []riskSummary) error {
	file, err := os.Create(*outFile)
	if err != nil {
		return errors.Wrap(err, "create CSV file")
	}
	defer file.Close()
	cw := csv.NewWriter(file)
	defer cw.Flush()

	if err := cw.Write([]string{"Cluster", "SignatureID", "Kind", "Name", "Namespace"}); err != nil {
		return errors.Wrap(err, "write CSV header")
	}

	for _, rs := range risks {
		resources, err := listRiskResources(ctx, token, rs.ID)
		if err != nil {
			return err
		}

		for _, r := range resources {
			if *namespace != "" && r.Namespace != *namespace {
				continue
			}
			if err := cw.Write([]string{
				*clusterID,
				rs.Signature.ID,
				r.Kind,
				r.Name,
				r.Namespace,
			}); err != nil {
				return errors.Wrap(err, "write CSV row")
			}
		}
	}
	return nil
}

func filterByNamespace(resources []riskResource, ns string) int {
	if ns == "" {
		return len(resources)
	}
	count := 0
	for _, r := range resources {
		if r.Namespace == ns {
			count++
		}
	}
	return count
}

func authenticate(ctx context.Context) (string, error) {
	url, err := presignSTS(ctx)
	if err != nil {
		return "", errors.Wrap(err, "generate presigned STS URL")
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiBase+"/login", nil)
	if err != nil {
		return "", errors.Wrap(err, "construct login request")
	}
	req.Header.Set("Authorization", "AWS "+url)

	resp, err := httpClient.Do(req)
	if err != nil {
		return "", errors.Wrap(err, "perform login request")
	}
	defer resp.Body.Close()

	var lr loginResponse
	if err := json.NewDecoder(resp.Body).Decode(&lr); err != nil {
		return "", errors.Wrap(err, "decode login response")
	}

	for _, acct := range lr.AccessTokens {
		for _, bundle := range acct {
			return bundle.AccessToken, nil
		}
	}
	return "", errors.New("no access tokens returned")
}

func listRisks(ctx context.Context, token, cluster string) ([]riskSummary, error) {
	url := fmt.Sprintf("%s/risks?filter=cluster_id:%s", apiBase, cluster)

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, errors.Wrap(err, "construct list risks request")
	}

	req.Header.Set("Authorization", "Bearer "+token)

	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, errors.Wrap(err, "perform list risks request")
	}
	defer resp.Body.Close()

	var lr listRisksResponse
	if err := json.NewDecoder(resp.Body).Decode(&lr); err != nil {
		return nil, errors.Wrap(err, "decode list risks response")
	}
	return lr.Data, nil
}

func listRiskResources(ctx context.Context, token, riskID string) ([]riskResource, error) {
	url := fmt.Sprintf("%s/risks/%s/resources", apiBase, riskID)

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, errors.Wrap(err, "construct list risk resources request")
	}
	req.Header.Set("Authorization", "Bearer "+token)

	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, errors.Wrap(err, "perform list risk resources request")
	}
	defer resp.Body.Close()

	var r listRiskResourcesResponse
	if err := json.NewDecoder(resp.Body).Decode(&r); err != nil {
		return nil, errors.Wrap(err, "decode list risk resources response")
	}
	return r.Data, nil
}

func presignSTS(ctx context.Context) (string, error) {
	cfg, err := config.LoadDefaultConfig(ctx)
	if err != nil {
		return "", errors.Wrap(err, "load AWS config")
	}
	stsClient := sts.NewFromConfig(cfg)
	sign := sts.NewPresignClient(stsClient)
	identity, err := sign.PresignGetCallerIdentity(ctx, &sts.GetCallerIdentityInput{})
	if err != nil {
		return "", errors.Wrap(err, "failed to create presigned URL")
	}
	return identity.URL, nil
}

func exitErr(err error) {
	fmt.Fprintf(os.Stderr, "ERROR: %+v\n", err)
	os.Exit(1)
}