From 9e3339b4f1b869b3aab626cc85fb5ebd6650ea3d Mon Sep 17 00:00:00 2001
From: Ward Vandewege <ward@jhvc.com>
Date: Sun, 17 Oct 2021 16:29:30 -0400
Subject: [PATCH] Add cli support for unsharing a node from a namespace.

---
 cmd/headscale/cli/nodes.go | 50 ++++++++++++++++++++++++++++++++++
 sharing.go                 | 24 ++++++++++++++++
 sharing_test.go            | 56 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 130 insertions(+)

diff --git a/cmd/headscale/cli/nodes.go b/cmd/headscale/cli/nodes.go
index 12461924..c44aa5ed 100644
--- a/cmd/headscale/cli/nodes.go
+++ b/cmd/headscale/cli/nodes.go
@@ -26,6 +26,7 @@ func init() {
 	nodeCmd.AddCommand(registerNodeCmd)
 	nodeCmd.AddCommand(deleteNodeCmd)
 	nodeCmd.AddCommand(shareMachineCmd)
+	nodeCmd.AddCommand(unshareMachineCmd)
 }
 
 var nodeCmd = &cobra.Command{
@@ -229,6 +230,55 @@ var shareMachineCmd = &cobra.Command{
 	},
 }
 
+var unshareMachineCmd = &cobra.Command{
+	Use:   "unshare ID",
+	Short: "Unshares a node from the specified namespace",
+	Args: func(cmd *cobra.Command, args []string) error {
+		if len(args) < 1 {
+			return fmt.Errorf("missing parameters")
+		}
+		return nil
+	},
+	Run: func(cmd *cobra.Command, args []string) {
+		namespace, err := cmd.Flags().GetString("namespace")
+		if err != nil {
+			log.Fatalf("Error getting namespace: %s", err)
+		}
+		output, _ := cmd.Flags().GetString("output")
+
+		h, err := getHeadscaleApp()
+		if err != nil {
+			log.Fatalf("Error initializing: %s", err)
+		}
+
+		n, err := h.GetNamespace(namespace)
+		if err != nil {
+			log.Fatalf("Error fetching namespace: %s", err)
+		}
+
+		id, err := strconv.Atoi(args[0])
+		if err != nil {
+			log.Fatalf("Error converting ID to integer: %s", err)
+		}
+		machine, err := h.GetMachineByID(uint64(id))
+		if err != nil {
+			log.Fatalf("Error getting node: %s", err)
+		}
+
+		err = h.RemoveSharedMachineFromNamespace(machine, n)
+		if strings.HasPrefix(output, "json") {
+			JsonOutput(map[string]string{"Result": "Node unshared"}, err, output)
+			return
+		}
+		if err != nil {
+			fmt.Printf("Error unsharing node: %s\n", err)
+			return
+		}
+
+		fmt.Println("Node unshared!")
+	},
+}
+
 func nodesToPtables(currentNamespace headscale.Namespace, machines []headscale.Machine) (pterm.TableData, error) {
 	d := pterm.TableData{{"ID", "Name", "NodeKey", "Namespace", "IP address", "Ephemeral", "Last seen", "Online"}}
 
diff --git a/sharing.go b/sharing.go
index 2905c8ef..879ed06f 100644
--- a/sharing.go
+++ b/sharing.go
@@ -40,6 +40,30 @@ func (h *Headscale) AddSharedMachineToNamespace(m *Machine, ns *Namespace) error
 	return nil
 }
 
+// RemoveSharedMachineFromNamespace removes a shared machine from a namespace
+func (h *Headscale) RemoveSharedMachineFromNamespace(m *Machine, ns *Namespace) error {
+	if m.NamespaceID == ns.ID {
+		return errorSameNamespace
+	}
+
+	sharedMachine := SharedMachine{}
+	result := h.db.Where("machine_id = ? AND namespace_id = ?", m.ID, ns.ID).Unscoped().Delete(&sharedMachine)
+	if result.Error != nil {
+		return result.Error
+	}
+
+	if result.RowsAffected == 0 {
+		return errorMachineNotShared
+	}
+
+	err := h.RequestMapUpdates(ns.ID)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // RemoveSharedMachineFromAllNamespaces removes a machine as a shared node from all namespaces
 func (h *Headscale) RemoveSharedMachineFromAllNamespaces(m *Machine) error {
 	sharedMachine := SharedMachine{}
diff --git a/sharing_test.go b/sharing_test.go
index d8cd8029..c6d26f68 100644
--- a/sharing_test.go
+++ b/sharing_test.go
@@ -4,6 +4,36 @@ import (
 	"gopkg.in/check.v1"
 )
 
+func CreateNodeNamespace(c *check.C, namespace, node, key, IP string) (*Namespace, *Machine) {
+	n1, err := h.CreateNamespace(namespace)
+	c.Assert(err, check.IsNil)
+
+	pak1, err := h.CreatePreAuthKey(n1.Name, false, false, nil)
+	c.Assert(err, check.IsNil)
+
+	_, err = h.GetMachine(n1.Name, node)
+	c.Assert(err, check.NotNil)
+
+	m1 := &Machine{
+		ID:             0,
+		MachineKey:     key,
+		NodeKey:        key,
+		DiscoKey:       key,
+		Name:           node,
+		NamespaceID:    n1.ID,
+		Registered:     true,
+		RegisterMethod: "authKey",
+		IPAddress:      IP,
+		AuthKeyID:      uint(pak1.ID),
+	}
+	h.db.Save(m1)
+
+	_, err = h.GetMachine(n1.Name, m1.Name)
+	c.Assert(err, check.IsNil)
+
+	return n1, m1
+}
+
 func (s *Suite) TestBasicSharedNodesInNamespace(c *check.C) {
 	n1, err := h.CreateNamespace("shared1")
 	c.Assert(err, check.IsNil)
@@ -125,6 +155,32 @@ func (s *Suite) TestSameNamespace(c *check.C) {
 	c.Assert(err, check.Equals, errorSameNamespace)
 }
 
+func (s *Suite) TestUnshare(c *check.C) {
+	n1, m1 := CreateNodeNamespace(c, "shared1", "test_unshare_1", "686824e749f3b7f2a5927ee6c1e422aee5292592d9179a271ed7b3e659b44a66", "100.64.0.1")
+	_, m2 := CreateNodeNamespace(c, "shared2", "test_unshare_2", "dec46ef9dc45c7d2f03bfcd5a640d9e24e3cc68ce3d9da223867c9bc6d5e9863", "100.64.0.2")
+
+	p1s, err := h.getPeers(m1)
+	c.Assert(err, check.IsNil)
+	c.Assert(len(p1s), check.Equals, 0)
+
+	err = h.AddSharedMachineToNamespace(m2, n1)
+	c.Assert(err, check.IsNil)
+
+	p1s, err = h.getShared(m1)
+	c.Assert(err, check.IsNil)
+	c.Assert(len(p1s), check.Equals, 1)
+
+	err = h.RemoveSharedMachineFromNamespace(m2, n1)
+	c.Assert(err, check.IsNil)
+
+	p1s, err = h.getShared(m1)
+	c.Assert(err, check.IsNil)
+	c.Assert(len(p1s), check.Equals, 0)
+
+	err = h.RemoveSharedMachineFromNamespace(m2, n1)
+	c.Assert(err, check.Equals, errorMachineNotShared)
+}
+
 func (s *Suite) TestAlreadyShared(c *check.C) {
 	n1, err := h.CreateNamespace("shared1")
 	c.Assert(err, check.IsNil)