11package premium
22
33import (
4+ "context"
45 "crypto/ed25519"
56 _ "embed"
67 "encoding/hex"
78 "encoding/json"
89 "errors"
10+ "fmt"
911 "os"
1012 "path/filepath"
1113 "slices"
1214 "strings"
1315 "time"
1416
17+ "github.com/aws/aws-sdk-go-v2/aws"
18+ awsConfig "github.com/aws/aws-sdk-go-v2/config"
19+ "github.com/aws/aws-sdk-go-v2/service/licensemanager"
20+ "github.com/aws/aws-sdk-go-v2/service/licensemanager/types"
1521 "github.com/cloudquery/plugin-sdk/v4/plugin"
22+ "github.com/google/uuid"
1623 "github.com/rs/zerolog"
1724)
1825
@@ -41,20 +48,88 @@ var publicKey string
4148
4249var timeFunc = time .Now
4350
44- func ValidateLicense (logger zerolog.Logger , meta plugin.Meta , licenseFileOrDirectory string ) error {
45- fi , err := os .Stat (licenseFileOrDirectory )
51+ const awsProductSKU = "prod-2trmtbe74klkg"
52+
53+ //go:generate mockgen -package=mocks -destination=../premium/mocks/licensemanager.go -source=offline.go AWSLicenseManagerInterface
54+ type AWSLicenseManagerInterface interface {
55+ CheckoutLicense (ctx context.Context , params * licensemanager.CheckoutLicenseInput , optFns ... func (* licensemanager.Options )) (* licensemanager.CheckoutLicenseOutput , error )
56+ }
57+
58+ type CQLicenseClient struct {
59+ logger zerolog.Logger
60+ meta plugin.Meta
61+ licenseFileOrDirectory string
62+ awsLicenseManagerClient AWSLicenseManagerInterface
63+ isMarketplaceLicense bool
64+ }
65+
66+ type LicenseClientOptions func (updater * CQLicenseClient )
67+
68+ func WithMeta (meta plugin.Meta ) LicenseClientOptions {
69+ return func (cl * CQLicenseClient ) {
70+ cl .meta = meta
71+ }
72+ }
73+
74+ func WithLicenseFileOrDirectory (licenseFileOrDirectory string ) LicenseClientOptions {
75+ return func (cl * CQLicenseClient ) {
76+ cl .licenseFileOrDirectory = licenseFileOrDirectory
77+ }
78+ }
79+
80+ func WithAWSLicenseManagerClient (awsLicenseManagerClient AWSLicenseManagerInterface ) LicenseClientOptions {
81+ return func (cl * CQLicenseClient ) {
82+ cl .awsLicenseManagerClient = awsLicenseManagerClient
83+ }
84+ }
85+
86+ func NewLicenseClient (ctx context.Context , logger zerolog.Logger , ops ... LicenseClientOptions ) (CQLicenseClient , error ) {
87+ cl := CQLicenseClient {
88+ logger : logger ,
89+ isMarketplaceLicense : os .Getenv ("CQ_AWS_MARKETPLACE_LICENSE" ) == "true" ,
90+ }
91+
92+ for _ , op := range ops {
93+ op (& cl )
94+ }
95+
96+ if cl .isMarketplaceLicense && cl .awsLicenseManagerClient == nil {
97+ cfg , err := awsConfig .LoadDefaultConfig (ctx )
98+ if err != nil {
99+ return cl , fmt .Errorf ("failed to load AWS config: %w" , err )
100+ }
101+ cl .awsLicenseManagerClient = licensemanager .NewFromConfig (cfg )
102+ }
103+
104+ return cl , nil
105+ }
106+
107+ func (lc CQLicenseClient ) ValidateLicense (ctx context.Context ) error {
108+ // License can be provided via environment variable for AWS Marketplace or CLI flag
109+ switch {
110+ case lc .isMarketplaceLicense :
111+ return lc .validateMarketplaceLicense (ctx )
112+ case lc .licenseFileOrDirectory != "" :
113+ return lc .validateCQLicense ()
114+ default :
115+ return ErrLicenseNotApplicable
116+ }
117+ }
118+
119+ func (lc CQLicenseClient ) validateCQLicense () error {
120+ fi , err := os .Stat (lc .licenseFileOrDirectory )
46121 if err != nil {
47122 return err
48123 }
49124 if ! fi .IsDir () {
50- return validateLicenseFile (logger , meta , licenseFileOrDirectory )
125+ return lc . validateLicenseFile (lc . licenseFileOrDirectory )
51126 }
52127
53128 found := false
54129 var lastError error
55- err = filepath .WalkDir (licenseFileOrDirectory , func (path string , d os.DirEntry , err error ) error {
130+ err = filepath .WalkDir (lc . licenseFileOrDirectory , func (path string , d os.DirEntry , err error ) error {
56131 if d .IsDir () {
57- if path == licenseFileOrDirectory {
132+ if path == lc . licenseFileOrDirectory {
58133 return nil
59134 }
60135 return filepath .SkipDir
@@ -67,8 +142,8 @@ func ValidateLicense(logger zerolog.Logger, meta plugin.Meta, licenseFileOrDirec
67142 return nil
68143 }
69144
70- logger .Debug ().Str ("path" , path ).Msg ("considering license file" )
71- lastError = validateLicenseFile (logger , meta , path )
145+ lc . logger .Debug ().Str ("path" , path ).Msg ("considering license file" )
146+ lastError = lc . validateLicenseFile (path )
72147 switch lastError {
73148 case nil :
74149 found = true
@@ -91,7 +166,7 @@ func ValidateLicense(logger zerolog.Logger, meta plugin.Meta, licenseFileOrDirec
91166 return errors .New ("failed to validate license directory" )
92167}
93168
94- func validateLicenseFile ( logger zerolog. Logger , meta plugin. Meta , licenseFile string ) error {
169+ func ( lc CQLicenseClient ) validateLicenseFile ( licenseFile string ) error {
95170 licenseContents , err := os .ReadFile (licenseFile )
96171 if err != nil {
97172 return err
@@ -103,14 +178,14 @@ func validateLicenseFile(logger zerolog.Logger, meta plugin.Meta, licenseFile st
103178 }
104179
105180 if len (l .Plugins ) > 0 {
106- ref := strings .Join ([]string {meta .Team , string (meta .Kind ), meta .Name }, "/" )
107- teamRef := meta .Team + "/*"
181+ ref := strings .Join ([]string {lc . meta .Team , string (lc . meta .Kind ), lc . meta .Name }, "/" )
182+ teamRef := lc . meta .Team + "/*"
108183 if ! slices .Contains (l .Plugins , ref ) && ! slices .Contains (l .Plugins , teamRef ) {
109184 return ErrLicenseNotApplicable
110185 }
111186 }
112187
113- return l .IsValid (logger )
188+ return l .IsValid (lc . logger )
114189}
115190
116191func UnpackLicense (lic []byte ) (* License , error ) {
@@ -158,3 +233,28 @@ func (l *License) IsValid(logger zerolog.Logger) error {
158233 msg .Time ("expires_at" , l .ExpiresAt ).Msgf ("Offline license for %s loaded." , l .LicensedTo )
159234 return nil
160235}
236+
237+ func (lc CQLicenseClient ) validateMarketplaceLicense (ctx context.Context ) error {
238+ clientToken := uuid .New ()
239+
240+ resp , err := lc .awsLicenseManagerClient .CheckoutLicense (ctx , & licensemanager.CheckoutLicenseInput {
241+ CheckoutType : types .CheckoutTypeProvisional ,
242+ ClientToken : aws .String (clientToken .String ()),
243+ ProductSKU : aws .String (awsProductSKU ),
244+ Entitlements : []types.EntitlementData {
245+ {
246+ Name : aws .String ("Unlimited" ),
247+ Unit : types .EntitlementDataUnitNone ,
248+ },
249+ },
250+ // This is hardcoded for AWS Marketplace, because this is the only supported value for marketplace licenses
251+ KeyFingerprint : aws .String ("aws:294406891311:AWS/Marketplace:issuer-fingerprint" ),
252+ })
253+ if err != nil {
254+ return fmt .Errorf ("failed to checkout license: %w" , err )
255+ }
256+ if len (resp .EntitlementsAllowed ) == 0 {
257+ return errors .New ("no entitlements provisioned" )
258+ }
259+ return nil
260+ }
0 commit comments