From f3139d26c87ca0bbcda014e9ecc6fe1734bcfab0 Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Sat, 17 Jul 2021 00:14:22 +0200
Subject: [PATCH] Added methods to delete nodes

---
 machine.go      | 25 +++++++++++++++++
 machine_test.go | 71 ++++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 95 insertions(+), 1 deletion(-)

diff --git a/machine.go b/machine.go
index cfc1b8f..5dcd571 100644
--- a/machine.go
+++ b/machine.go
@@ -188,6 +188,31 @@ func (h *Headscale) GetMachine(namespace string, name string) (*Machine, error)
 	return nil, fmt.Errorf("not found")
 }
 
+// GetMachineByID finds a Machine by ID and returns the Machine struct
+func (h *Headscale) GetMachineByID(id uint64) (*Machine, error) {
+	m := Machine{}
+	if result := h.db.Find(&Machine{ID: id}).First(&m); result.Error != nil {
+		return nil, result.Error
+	}
+	return &m, nil
+}
+
+// DeleteMachine softs deletes a Machine from the database
+func (h *Headscale) DeleteMachine(m Machine) error {
+	if err := h.db.Delete(&m).Error; err != nil {
+		return err
+	}
+	return nil
+}
+
+// HardDeleteMachine hard deletes a Machine from the database
+func (h *Headscale) HardDeleteMachine(m Machine) error {
+	if err := h.db.Unscoped().Delete(&m).Error; err != nil {
+		return err
+	}
+	return nil
+}
+
 // GetHostInfo returns a Hostinfo struct for the machine
 func (m *Machine) GetHostInfo() (*tailcfg.Hostinfo, error) {
 	hostinfo := tailcfg.Hostinfo{}
diff --git a/machine_test.go b/machine_test.go
index d9a472c..2e6ae1f 100644
--- a/machine_test.go
+++ b/machine_test.go
@@ -32,5 +32,74 @@ func (s *Suite) TestGetMachine(c *check.C) {
 
 	_, err = m1.GetHostInfo()
 	c.Assert(err, check.IsNil)
-
+}
+
+func (s *Suite) TestGetMachineByID(c *check.C) {
+	n, err := h.CreateNamespace("test")
+	c.Assert(err, check.IsNil)
+
+	pak, err := h.CreatePreAuthKey(n.Name, false, false, nil)
+	c.Assert(err, check.IsNil)
+
+	_, err = h.GetMachineByID(0)
+	c.Assert(err, check.NotNil)
+
+	m := Machine{
+		ID:             0,
+		MachineKey:     "foo",
+		NodeKey:        "bar",
+		DiscoKey:       "faa",
+		Name:           "testmachine",
+		NamespaceID:    n.ID,
+		Registered:     true,
+		RegisterMethod: "authKey",
+		AuthKeyID:      uint(pak.ID),
+	}
+	h.db.Save(&m)
+
+	m1, err := h.GetMachineByID(0)
+	c.Assert(err, check.IsNil)
+
+	_, err = m1.GetHostInfo()
+	c.Assert(err, check.IsNil)
+}
+
+func (s *Suite) TestDeleteMachine(c *check.C) {
+	n, err := h.CreateNamespace("test")
+	c.Assert(err, check.IsNil)
+	m := Machine{
+		ID:             0,
+		MachineKey:     "foo",
+		NodeKey:        "bar",
+		DiscoKey:       "faa",
+		Name:           "testmachine",
+		NamespaceID:    n.ID,
+		Registered:     true,
+		RegisterMethod: "authKey",
+		AuthKeyID:      uint(1),
+	}
+	h.db.Save(&m)
+	h.DeleteMachine(m)
+	_, err = h.GetMachine(n.Name, "testmachine")
+	c.Assert(err, check.NotNil)
+}
+
+func (s *Suite) TestHardDeleteMachine(c *check.C) {
+	n, err := h.CreateNamespace("test")
+	c.Assert(err, check.IsNil)
+	m := Machine{
+		ID:             0,
+		MachineKey:     "foo",
+		NodeKey:        "bar",
+		DiscoKey:       "faa",
+		Name:           "testmachine3",
+		NamespaceID:    n.ID,
+		Registered:     true,
+		RegisterMethod: "authKey",
+		AuthKeyID:      uint(1),
+	}
+	h.db.Save(&m)
+	h.HardDeleteMachine(m)
+	_, err = h.GetMachine(n.Name, "testmachine3")
+	c.Assert(err, check.NotNil)
 }