bitbucket.org/Aishee/synsec@v0.0.0-20210414005726-236fc01a153d/pkg/database/machines.go (about)

     1  package database
     2  
     3  import (
     4  	"fmt"
     5  	"time"
     6  
     7  	"github.com/go-openapi/strfmt"
     8  
     9  	"bitbucket.org/Aishee/synsec/pkg/database/ent"
    10  	"bitbucket.org/Aishee/synsec/pkg/database/ent/machine"
    11  	"github.com/pkg/errors"
    12  	"golang.org/x/crypto/bcrypt"
    13  )
    14  
    15  func (c *Client) CreateMachine(machineID *string, password *strfmt.Password, ipAddress string, isValidated bool, force bool) (int, error) {
    16  	hashPassword, err := bcrypt.GenerateFromPassword([]byte(*password), bcrypt.DefaultCost)
    17  	if err != nil {
    18  		c.Log.Warningf("CreateMachine : %s", err)
    19  		return 0, errors.Wrap(HashError, "")
    20  	}
    21  
    22  	machineExist, err := c.Ent.Machine.
    23  		Query().
    24  		Where(machine.MachineIdEQ(*machineID)).
    25  		Select(machine.FieldMachineId).Strings(c.CTX)
    26  	if err != nil {
    27  		return 0, errors.Wrapf(QueryFail, "machine '%s': %s", *machineID, err)
    28  	}
    29  	if len(machineExist) > 0 {
    30  		if force {
    31  			_, err := c.Ent.Machine.Update().Where(machine.MachineIdEQ(*machineID)).SetPassword(string(hashPassword)).Save(c.CTX)
    32  			if err != nil {
    33  				c.Log.Warningf("CreateMachine : %s", err)
    34  				return 0, errors.Wrapf(UpdateFail, "machine '%s'", *machineID)
    35  			}
    36  			return 1, nil
    37  		}
    38  		return 0, errors.Wrapf(UserExists, "user '%s'", *machineID)
    39  	}
    40  
    41  	_, err = c.Ent.Machine.
    42  		Create().
    43  		SetMachineId(*machineID).
    44  		SetPassword(string(hashPassword)).
    45  		SetIpAddress(ipAddress).
    46  		SetIsValidated(isValidated).
    47  		Save(c.CTX)
    48  
    49  	if err != nil {
    50  		c.Log.Warningf("CreateMachine : %s", err)
    51  		return 0, errors.Wrapf(InsertFail, "creating machine '%s'", *machineID)
    52  	}
    53  
    54  	return 1, nil
    55  }
    56  
    57  func (c *Client) QueryMachineByID(machineID string) (*ent.Machine, error) {
    58  	machine, err := c.Ent.Machine.
    59  		Query().
    60  		Where(machine.MachineIdEQ(machineID)).
    61  		Only(c.CTX)
    62  	if err != nil {
    63  		c.Log.Warningf("QueryMachineByID : %s", err)
    64  		return &ent.Machine{}, errors.Wrapf(UserNotExists, "user '%s'", machineID)
    65  	}
    66  	return machine, nil
    67  }
    68  
    69  func (c *Client) ListMachines() ([]*ent.Machine, error) {
    70  	machines, err := c.Ent.Machine.Query().All(c.CTX)
    71  	if err != nil {
    72  		return []*ent.Machine{}, errors.Wrapf(QueryFail, "listing machines: %s", err)
    73  	}
    74  	return machines, nil
    75  }
    76  
    77  func (c *Client) ValidateMachine(machineID string) error {
    78  	_, err := c.Ent.Machine.Update().Where(machine.MachineIdEQ(machineID)).SetIsValidated(true).Save(c.CTX)
    79  	if err != nil {
    80  		return errors.Wrapf(UpdateFail, "validating machine: %s", err)
    81  	}
    82  	return nil
    83  }
    84  
    85  func (c *Client) QueryPendingMachine() ([]*ent.Machine, error) {
    86  	var machines []*ent.Machine
    87  	var err error
    88  
    89  	machines, err = c.Ent.Machine.Query().Where(machine.IsValidatedEQ(false)).All(c.CTX)
    90  	if err != nil {
    91  		c.Log.Warningf("QueryPendingMachine : %s", err)
    92  		return []*ent.Machine{}, errors.Wrapf(QueryFail, "querying pending machines: %s", err)
    93  	}
    94  	return machines, nil
    95  }
    96  
    97  func (c *Client) DeleteWatcher(name string) error {
    98  	_, err := c.Ent.Machine.
    99  		Delete().
   100  		Where(machine.MachineIdEQ(name)).
   101  		Exec(c.CTX)
   102  	if err != nil {
   103  		return fmt.Errorf("unable to save api key in database: %s", err)
   104  	}
   105  	return nil
   106  }
   107  
   108  func (c *Client) UpdateMachineScenarios(scenarios string, ID int) error {
   109  	_, err := c.Ent.Machine.UpdateOneID(ID).
   110  		SetUpdatedAt(time.Now()).
   111  		SetScenarios(scenarios).
   112  		Save(c.CTX)
   113  	if err != nil {
   114  		return fmt.Errorf("unable to update machine in database: %s", err)
   115  	}
   116  	return nil
   117  }
   118  
   119  func (c *Client) UpdateMachineIP(ipAddr string, ID int) error {
   120  	_, err := c.Ent.Machine.UpdateOneID(ID).
   121  		SetIpAddress(ipAddr).
   122  		Save(c.CTX)
   123  	if err != nil {
   124  		return fmt.Errorf("unable to update machine in database: %s", err)
   125  	}
   126  	return nil
   127  }
   128  
   129  func (c *Client) UpdateMachineVersion(ipAddr string, ID int) error {
   130  	_, err := c.Ent.Machine.UpdateOneID(ID).
   131  		SetVersion(ipAddr).
   132  		Save(c.CTX)
   133  	if err != nil {
   134  		return fmt.Errorf("unable to update machine in database: %s", err)
   135  	}
   136  	return nil
   137  }
   138  
   139  func (c *Client) IsMachineRegistered(machineID string) (bool, error) {
   140  	exist, err := c.Ent.Machine.Query().Where().Select(machine.FieldMachineId).Strings(c.CTX)
   141  	if err != nil {
   142  		return false, err
   143  	}
   144  	if len(exist) == 1 {
   145  		return true, nil
   146  	}
   147  	if len(exist) > 1 {
   148  		return false, fmt.Errorf("More than one item with the same machineID in database")
   149  	}
   150  
   151  	return false, nil
   152  
   153  }