timeout support added

appknox · Nov 21, 2024 · eaa3f3a · eaa3f3a
1 parent 3b69b96
commit eaa3f3a
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 6 deletions.
diff --git a/appknox/sarif_generator.go b/appknox/sarif_generator.go
@@ -13,6 +13,7 @@ import (
 	"github.com/iancoleman/strcase"
 	"github.com/vbauerster/mpb/v4"
 	"github.com/vbauerster/mpb/v4/decor"
+	"github.com/spf13/viper"
 )
 
 type SARIF struct {
@@ -109,7 +110,7 @@ func GenerateSARIFGivenFileID(client *Client, fileID int, riskThreshold int) (SA
 			decor.Name("] "),
 		),
 	)
-
+	
 	for sarifReportProgess < 100 {
 		file, _, err := client.Files.GetByID(ctx, fileID)
 		if err != nil {
@@ -118,7 +119,8 @@ func GenerateSARIFGivenFileID(client *Client, fileID int, riskThreshold int) (SA
 		}
 		sarifReportProgess = file.StaticScanProgress
 		bar.SetCurrent(int64(sarifReportProgess), time.Since(start))
-		if time.Since(start) > 15*time.Minute {
+
+		if time.Since(start) > time.Duration(viper.GetInt("timeout")) * time.Minute {
 			err := errors.New("Request timed out")
 			PrintError(err)
 			os.Exit(1)

diff --git a/cmd/root.go b/cmd/root.go
@@ -4,7 +4,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
-
+	
 	// "github.com/appknox/appknox-go/appknox"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -39,7 +39,11 @@ func init() {
     viper.BindPFlag("host", RootCmd.PersistentFlags().Lookup("host"))
     viper.BindEnv("host", "APPKNOX_API_HOST")
 
-
+	// Define flags globally here for all subcommands
+	RootCmd.PersistentFlags().String("timeout", "", "Timeout for Appknox scanner, default is 30 minutes")
+	viper.BindPFlag("timeout", RootCmd.PersistentFlags().Lookup("timeout"))
+	viper.BindEnv("timeout", "APPKNOX_TIMEOUT")
+	viper.SetDefault("timeout", "30") // Default to 30 minutes
 	RootCmd.PersistentFlags().String("region", "", "Region names, e.g., global, saudi, uae. By default, global is used")
     viper.BindPFlag("region", RootCmd.PersistentFlags().Lookup("region"))
     viper.BindEnv("region", "APPKNOX_API_REGION")

diff --git a/helper/cicheck.go b/helper/cicheck.go
@@ -6,9 +6,10 @@ import (
 	"fmt"
 	"os"
 	"time"
-
+	
 	"github.com/appknox/appknox-go/appknox"
 	"github.com/appknox/appknox-go/appknox/enums"
+	"github.com/spf13/viper"
 	"github.com/cheynewallace/tabby"
 	"github.com/vbauerster/mpb/v4"
 	"github.com/vbauerster/mpb/v4/decor"
@@ -20,6 +21,9 @@ func ProcessCiCheck(fileID, riskThreshold int) {
 	client := getClient()
 	var staticScanProgess int
 	start := time.Now()
+	timeoutMinutes := viper.GetInt("timeout")
+	timeout := time.Duration(timeoutMinutes) * time.Minute
+	fmt.Printf("Starting scan at: %v with timeout of %d minutes\n", start.Format(time.RFC3339), timeoutMinutes)
 	p := mpb.New(
 		mpb.WithWidth(60),
 		mpb.WithRefreshRate(180*time.Millisecond),
@@ -44,7 +48,8 @@ func ProcessCiCheck(fileID, riskThreshold int) {
 		}
 		staticScanProgess = file.StaticScanProgress
 		bar.SetCurrent(int64(staticScanProgess), time.Since(start))
-		if time.Since(start) > 30*time.Minute {
+
+		if time.Since(start) > timeout {
 			err := errors.New("Request timed out")
 			PrintError(err)
 			os.Exit(1)