Merge pull request #5 from bytedance/use-internalip

Use InternalIP to connect to nodes.
sysdiglabs · Oct 17, 2022 · f525558 · f525558
2 parents 6b8ff07 + fe88c63
commit f525558
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 6 deletions.
diff --git a/aa/service.go b/aa/service.go
@@ -25,6 +25,7 @@ const (
 type AppArmor struct {
 	k8sClient *client.K8sClient
 	sshClient *client.SSHClient
+	useInternalIP bool
 }
 
 // NewAppArmor returns a new AppArmor object
@@ -57,9 +58,14 @@ func NewAppArmor() (*AppArmor, error) {
 	return &AppArmor{
 		k8sClient: k8s,
 		sshClient: ssh,
+		useInternalIP: false,
 	}, nil
 }
 
+func (aa *AppArmor) UseInternalIP(useInternalIP bool) {
+	aa.useInternalIP = useInternalIP
+}
+
 // InstallCRD installs CRD in Kubernetes
 func (aa *AppArmor) InstallCRD() error {
 	return aa.k8sClient.InstallCRD()
@@ -88,7 +94,12 @@ func (aa *AppArmor) install(node *types.Node) error {
 		return nil
 	}
 
-	err := aa.sshClient.Connect(node.ExternalIP, SSH_PORT)
+	var err error
+	if aa.useInternalIP {
+		err = aa.sshClient.Connect(node.InternalIP, SSH_PORT)
+	} else {
+		err = aa.sshClient.Connect(node.ExternalIP, SSH_PORT)
+	}
 
 	if err != nil {
 		return err
@@ -97,7 +108,11 @@ func (aa *AppArmor) install(node *types.Node) error {
 	defer aa.sshClient.Close()
 
 	if aa.enabledInConnection(node) {
-		klog.Infof("AppArmor was enabled on node: %s (external IP: %s)", node.NodeName, node.ExternalIP)
+		if aa.useInternalIP {
+			klog.Infof("AppArmor was enabled on node: %s (internal IP: %s)", node.NodeName, node.InternalIP)
+		} else {
+			klog.Infof("AppArmor was enabled on node: %s (external IP: %s)", node.NodeName, node.ExternalIP)
+		}
 		return nil
 	}
 
@@ -141,15 +156,25 @@ func (aa *AppArmor) syncProfile(node *types.Node, profile types.AppArmorProfile)
 		return nil
 	}
 
-	err := aa.sshClient.Connect(node.ExternalIP, SSH_PORT)
+	var err error
+	if aa.useInternalIP {
+		err = aa.sshClient.Connect(node.InternalIP, SSH_PORT)
+	} else {
+		err = aa.sshClient.Connect(node.ExternalIP, SSH_PORT)
+	}
+
 	if err != nil {
 		return err
 	}
 
 	defer aa.sshClient.Close()
 
 	if !aa.enabledInConnection(node) {
-		klog.Infof("AppArmor was not enabled on node: %s (external IP: %s), no sync happen.", node.NodeName, node.ExternalIP)
+		if aa.useInternalIP {
+			klog.Infof("AppArmor was not enabled on node: %s (internal IP: %s), no sync happen.", node.NodeName, node.InternalIP)
+		} else {
+			klog.Infof("AppArmor was not enabled on node: %s (external IP: %s), no sync happen.", node.NodeName, node.ExternalIP)
+		}
 		return nil
 	}
 
@@ -196,7 +221,12 @@ func (aa *AppArmor) enabled(node *types.Node) (bool, error) {
 		return false, nil
 	}
 
-	err := aa.sshClient.Connect(node.ExternalIP, SSH_PORT)
+	var err error
+	if aa.useInternalIP {
+		err = aa.sshClient.Connect(node.InternalIP, SSH_PORT)
+	} else {
+		err = aa.sshClient.Connect(node.ExternalIP, SSH_PORT)
+	}
 	if err != nil {
 		return false, err
 	}
@@ -249,7 +279,12 @@ func (aa *AppArmor) status(node *types.Node) error {
 		return nil
 	}
 
-	err := aa.sshClient.Connect(node.ExternalIP, SSH_PORT)
+	var err error
+	if aa.useInternalIP {
+		err = aa.sshClient.Connect(node.InternalIP, SSH_PORT)
+	} else {
+		err = aa.sshClient.Connect(node.ExternalIP, SSH_PORT)
+	}
 	if err != nil {
 		return err
 	}

diff --git a/main.go b/main.go
@@ -20,6 +20,7 @@ func main() {
 	}
 
 	var logLevel string
+	var useInternalIP bool
 
 	log.SetFormatter(&log.TextFormatter{
 		FullTimestamp: true,
@@ -44,10 +45,12 @@ func main() {
 			}
 
 			log.SetLevel(lvl)
+			appArmor.UseInternalIP(useInternalIP)
 		},
 	}
 
 	rootCmd.PersistentFlags().StringVar(&logLevel, "level", "info", "Log level")
+	rootCmd.PersistentFlags().BoolVarP(&useInternalIP, "internal-ip", "i", false, "Use internal ip to sync")
 
 	var initCmd = &cobra.Command{
 		Use:   "init",