diff --git a/machine.go b/machine.go index cfc1b8f8..5dcd5714 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 d9a472cb..2e6ae1f6 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) }