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.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()
if resp.StatusCode != 200 {
data, err := io.ReadAll(resp.Body)
if err != nil {
return "", errors.Wrap(err, "failed to read response body")
}
return "", errors.Errorf("Received Login Error. Code: %d Body: %s", resp.StatusCode, string(data))
}
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")
}
cfg.Region = "us-east-1"
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)
}
