final touches

cmd/policy-assistant/pkg/cli/analyze.go

@@ -42,7 +42,7 @@ var AllModes = []string{
 	ProbeMode,
 }
 
-const DefaultTimeout = 180
+const DefaultTimeout = 3 * time.Minute
 
 type AnalyzeArgs struct {
 	AllNamespaces      bool
@@ -63,7 +63,7 @@ type AnalyzeArgs struct {
 	// synthetic probe
 	ProbePath string
 
-	Timeout int
+	Timeout time.Duration
 }
 
 func SetupAnalyzeCommand() *cobra.Command {
@@ -90,8 +90,7 @@ func SetupAnalyzeCommand() *cobra.Command {
 	command.Flags().StringVar(&args.TargetPodPath, "target-pod-path", "", "path to json target pod file -- json array of dicts")
 	command.Flags().StringVar(&args.TrafficPath, "traffic-path", "", "path to json traffic file, containing of a list of traffic objects")
 	command.Flags().StringVar(&args.ProbePath, "probe-path", "", "path to json model file for synthetic probe")
-
-	command.Flags().IntVar(&args.Timeout, "timeout", DefaultTimeout, "timeout time in seconds")
+	command.Flags().DurationVar(&args.Timeout, "kube-client-timeout", DefaultTimeout, "kube client timeout")
 
 	return command
 }
@@ -118,7 +117,7 @@ func RunAnalyzeCommand(args *AnalyzeArgs) {
 
 		includeANPS, includeBANPSs := shouldIncludeANPandBANP(kubeClient.ClientSet)
 
-		ctx, cancel := context.WithTimeout(context.TODO(), time.Duration(args.Timeout)*time.Second)
+		ctx, cancel := context.WithTimeout(context.TODO(), args.Timeout)
 		defer cancel()
 
 		kubePolicies, kubeANPs, kubeBANP, netpolErr, anpErr, banpErr = kube.ReadNetworkPoliciesFromKube(ctx, kubeClient, namespaces, includeANPS, includeBANPSs)
@@ -349,7 +348,7 @@ func shouldIncludeANPandBANP(client *kubernetes.Clientset) (bool, bool) {
 			case "BaselineAdminNetworkPolicy":
 				includeBANP = true
 			default:
-				panic(fmt.Sprintf("unexpected resource kind %s", res.Kind))
+				continue
 			}
 		}
 	}

cmd/policy-assistant/pkg/kube/ikubernetes.go

@@ -35,7 +35,7 @@ type IKubernetes interface {
 	UpdateAdminNetworkPolicy(ctx context.Context, policy *v1alpha1.AdminNetworkPolicy) (*v1alpha1.AdminNetworkPolicy, error)
 	DeleteAdminNetworkPolicy(ctx context.Context, name string) error
 
-	GetBaseAdminNetworkPolicies(ctx context.Context) ([]v1alpha1.BaselineAdminNetworkPolicy, error)
+	GetBaselineAdminNetworkPolicies(ctx context.Context) ([]v1alpha1.BaselineAdminNetworkPolicy, error)
 	CreateBaselineAdminNetworkPolicy(ctx context.Context, policy *v1alpha1.BaselineAdminNetworkPolicy) (*v1alpha1.BaselineAdminNetworkPolicy, error)
 	UpdateBaselineAdminNetworkPolicy(ctx context.Context, policy *v1alpha1.BaselineAdminNetworkPolicy) (*v1alpha1.BaselineAdminNetworkPolicy, error)
 	DeleteBaselineAdminNetworkPolicy(ctx context.Context, name string) error
@@ -95,12 +95,24 @@ func GetServicesInNamespaces(kubernetes IKubernetes, namespaces []string) ([]v1.
 	return allServices, nil
 }
 
-func GetAdminNetworkPoliciesInNamespaces(ctx context.Context, kubernetes IKubernetes) ([]v1alpha1.AdminNetworkPolicy, error) {
+func GetAdminNetworkPolicies(ctx context.Context, kubernetes IKubernetes) ([]v1alpha1.AdminNetworkPolicy, error) {
 	return kubernetes.GetAdminNetworkPolicies(ctx)
 }
 
-func GetBaseAdminNetworkPoliciesInNamespaces(ctx context.Context, kubernetes IKubernetes) ([]v1alpha1.BaselineAdminNetworkPolicy, error) {
-	return kubernetes.GetBaseAdminNetworkPolicies(ctx)
+func GetBaselineAdminNetworkPolicies(ctx context.Context, kubernetes IKubernetes) (*v1alpha1.BaselineAdminNetworkPolicy, error) {
+	result, err := kubernetes.GetBaselineAdminNetworkPolicies(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(result) == 0 {
+		return nil, nil
+	}
+
+	if len(result) == 1 {
+		return &result[0], nil
+	}
+	panic("Only one Banp should be returned")
 }
 
 type MockNamespace struct {
@@ -113,7 +125,7 @@ type MockNamespace struct {
 type MockKubernetes struct {
 	AdminNetworkPolicies        []v1alpha1.AdminNetworkPolicy
 	AdminNetworkPolicyError     error
-	BaseNetworkPolicies         []v1alpha1.BaselineAdminNetworkPolicy
+	BaselineNetworkPolicies     []v1alpha1.BaselineAdminNetworkPolicy
 	BaseAdminNetworkPolicyError error
 	Namespaces                  map[string]*MockNamespace
 	NetworkPolicyError          error
@@ -410,8 +422,8 @@ func (k *MockKubernetes) DeleteAdminNetworkPolicy(ctx context.Context, name stri
 	return ErrNotImplemented
 }
 
-func (m *MockKubernetes) GetBaseAdminNetworkPolicies(ctx context.Context) ([]v1alpha1.BaselineAdminNetworkPolicy, error) {
-	return m.BaseNetworkPolicies, m.BaseAdminNetworkPolicyError
+func (m *MockKubernetes) GetBaselineAdminNetworkPolicies(ctx context.Context) ([]v1alpha1.BaselineAdminNetworkPolicy, error) {
+	return m.BaselineNetworkPolicies, m.BaseAdminNetworkPolicyError
 }
 
 func (k *MockKubernetes) CreateBaselineAdminNetworkPolicy(ctx context.Context, policy *v1alpha1.BaselineAdminNetworkPolicy) (*v1alpha1.BaselineAdminNetworkPolicy, error) {

cmd/policy-assistant/pkg/kube/kubernetes.go

@@ -132,7 +132,7 @@ func (k *Kubernetes) DeleteAdminNetworkPolicy(ctx context.Context, name string)
 	return ErrNotImplemented
 }
 
-func (k *Kubernetes) GetBaseAdminNetworkPolicies(ctx context.Context) ([]v1alpha12.BaselineAdminNetworkPolicy, error) {
+func (k *Kubernetes) GetBaselineAdminNetworkPolicies(ctx context.Context) ([]v1alpha12.BaselineAdminNetworkPolicy, error) {
 	banp, err := k.alphaClientSet.BaselineAdminNetworkPolicies().List(ctx, metav1.ListOptions{})
 	if err != nil {
 		return nil, err

cmd/policy-assistant/pkg/kube/read.go

@@ -136,7 +136,7 @@ func ReadNetworkPoliciesFromKube(ctx context.Context, kubeClient IKubernetes, na
 		if !includeANPs {
 			return
 		}
-		anps, anpErr = GetAdminNetworkPoliciesInNamespaces(ctx, kubeClient)
+		anps, anpErr = GetAdminNetworkPolicies(ctx, kubeClient)
 		return
 	}(&wg)
 
@@ -145,14 +145,7 @@ func ReadNetworkPoliciesFromKube(ctx context.Context, kubeClient IKubernetes, na
 		if !includeBANPs {
 			return
 		}
-		result, err := GetBaseAdminNetworkPoliciesInNamespaces(ctx, kubeClient)
-		if err != nil {
-			banpErr = err
-		}
-		if len(result) > 0 {
-			banp = &result[0]
-		}
-
+		banp, banpErr = GetBaselineAdminNetworkPolicies(ctx, kubeClient)
 		return
 	}(&wg)
 

cmd/policy-assistant/pkg/kube/read_test.go

@@ -61,7 +61,7 @@ func TestReadNetworkPoliciesFromKube(t *testing.T) {
 			k := &MockKubernetes{
 				AdminNetworkPolicies:        scenario.AdminNetworkPolicies,
 				AdminNetworkPolicyError:     scenario.expectedAnpErr,
-				BaseNetworkPolicies:         scenario.BaselineAdminNetworkPolicies,
+				BaselineNetworkPolicies:     scenario.BaselineAdminNetworkPolicies,
 				BaseAdminNetworkPolicyError: scenario.expectedBanpErr,
 				Namespaces:                  map[string]*MockNamespace{},
 				NetworkPolicyError:          scenario.expectedNetErr,