diff --git a/api/client.go b/api/client.go index 540ab1b..aa52081 100644 --- a/api/client.go +++ b/api/client.go @@ -9,7 +9,7 @@ import ( "golang.org/x/oauth2/clientcredentials" ) -const DefaultTimeout = 5 +const DefaultTimeout = 5 * time.Second type Client struct { AuthConfig *clientcredentials.Config @@ -18,14 +18,18 @@ type Client struct { AuthToken string } -func NewClient(url, token string) (c *Client) { +func NewClient(url, token string, timeout time.Duration) (c *Client) { c = &Client{ ManagementURL: url, AuthToken: token, } + if timeout <= 0 { + timeout = DefaultTimeout + } + c.HTTPClient = NewLoggingHTTPClient() - c.HTTPClient.Timeout = time.Duration(DefaultTimeout) * time.Second + c.HTTPClient.Timeout = timeout return } diff --git a/api/client_test.go b/api/client_test.go index fbedbfd..0c5aac8 100644 --- a/api/client_test.go +++ b/api/client_test.go @@ -16,13 +16,13 @@ func envClient(t *testing.T) *api.Client { t.Skip("SENTINELONE_URL and SENTINELONE_TOKEN must be set!") } - return api.NewClient(url, token) + return api.NewClient(url, token, 0) } func testClient() (*api.Client, func()) { httpmock.Activate() - return api.NewClient("https://euce1-test.sentinelone.net", "test"), func() { + return api.NewClient("https://euce1-test.sentinelone.net", "test", 0), func() { httpmock.DeactivateAndReset() } } diff --git a/check.go b/check.go index 0691cd1..a106d0a 100644 --- a/check.go +++ b/check.go @@ -6,6 +6,7 @@ import ( "net/url" "os" "strings" + "time" "github.com/NETWAYS/check_sentinelone/api" "github.com/NETWAYS/go-check" @@ -19,6 +20,7 @@ type Config struct { IgnoreInProgress bool SiteName string ComputerName string + Timeout time.Duration } func BuildConfigFlags(fs *pflag.FlagSet) (config *Config) { @@ -58,7 +60,7 @@ func (c *Config) Validate() error { } func (c *Config) Run() (*result.PartialResult, error) { - client := api.NewClient(c.ManagementURL, c.AuthToken) + client := api.NewClient(c.ManagementURL, c.AuthToken, c.Timeout) values := url.Values{} values.Set("sortOrder", "desc") diff --git a/main.go b/main.go index 66bbd32..ba63667 100644 --- a/main.go +++ b/main.go @@ -1,6 +1,8 @@ package main import ( + "time" + "github.com/NETWAYS/go-check" "github.com/NETWAYS/go-check/result" ) @@ -27,6 +29,8 @@ func main() { plugin.ParseArguments() config.SetFromEnv() + config.Timeout = time.Duration(plugin.Timeout) * time.Second + err := config.Validate() if err != nil { check.ExitError(err)