move implantbuild models to protobuf struct

protobuf/clientpb/client.proto

@@ -241,6 +241,20 @@ message ExternalImplantBinary {
 // Configs of previously built implants
 message ImplantBuilds { map<string, ImplantConfig> Configs = 1; }
 
+message ImplantBuild {
+  string ID = 1;
+  string Name = 2;
+
+	string MD5    = 3;
+	string SHA1   = 4;
+	string SHA256 = 5;
+
+  bool Burned = 6;
+	uint64 ImplantID = 7;
+
+	string ImplantConfigID = 8;
+}
+
 message CompilerTarget {
   string GOOS = 1;   // The server's OS
   string GOARCH = 2; // The server's Arch

server/db/helpers.go

@@ -125,8 +125,24 @@ func ImplantBuilds() (*clientpb.ImplantBuilds, error) {
 	return pbBuilds, err
 }
 
+// SaveImplantBuild
+func SaveImplantBuild(ib *clientpb.ImplantBuild) (*clientpb.ImplantBuild, error) {
+	implantBuild := models.ImplantBuildFromProtobuf(ib)
+	dbSession := Session()
+	err := dbSession.Create(&implantBuild).Error
+	if err != nil {
+		return nil, err
+	}
+	build, err := ImplantBuildByName(implantBuild.Name)
+	if err != nil {
+		return nil, err
+	}
+
+	return build, nil
+}
+
 // ImplantBuildByName - Fetch implant build by name
-func ImplantBuildByName(name string) (*models.ImplantBuild, error) {
+func ImplantBuildByName(name string) (*clientpb.ImplantBuild, error) {
 	if len(name) < 1 {
 		return nil, ErrRecordNotFound
 	}
@@ -137,33 +153,19 @@ func ImplantBuildByName(name string) (*models.ImplantBuild, error) {
 	if err != nil {
 		return nil, err
 	}
-	return &build, err
-}
-
-// ImplantBuildNames - Fetch a list of all build names
-func ImplantBuildNames() ([]string, error) {
-	builds := []*models.ImplantBuild{}
-	err := Session().Where(&models.ImplantBuild{}).Find(&builds).Error
-	if err != nil {
-		return []string{}, err
-	}
-	names := []string{}
-	for _, build := range builds {
-		names = append(names, build.Name)
-	}
-	return names, nil
+	return build.ToProtobuf(), err
 }
 
 // ImplantBuildByResourceID - Fetch implant build from resource ID
-func ImplantBuildByResourceID(resourceID uint64) (*models.ImplantBuild, error) {
+func ImplantBuildByResourceID(resourceID uint64) (*clientpb.ImplantBuild, error) {
 	build := models.ImplantBuild{}
 	err := Session().Where(&models.ImplantBuild{
 		ImplantID: resourceID,
 	}).Find(&build).Error
 	if err != nil {
 		return nil, err
 	}
-	return &build, nil
+	return build.ToProtobuf(), nil
 }
 
 // ImplantProfiles - Fetch a map of name<->profiles current in the database
@@ -749,14 +751,17 @@ func BeaconTaskByID(id string) (*clientpb.BeaconTask, error) {
 	if len(id) < 1 {
 		return nil, ErrRecordNotFound
 	}
-	taskID := uuid.FromStringOrNil(id)
+	taskID, err := uuid.FromString(id)
 	if taskID == uuid.Nil {
 		return nil, ErrRecordNotFound
 	}
 	task := &models.BeaconTask{}
-	err := Session().Where(
+	err = Session().Where(
 		&models.BeaconTask{ID: taskID},
 	).First(task).Error
+	if err != nil {
+		return nil, err
+	}
 	return task.ToProtobuf(true), err
 }
 

server/db/models/beacon.go

@@ -124,8 +124,8 @@ type BeaconTask struct {
 	BeaconID    uuid.UUID `gorm:"type:uuid;"`
 	CreatedAt   time.Time `gorm:"->;<-:create;"`
 	State       string
-	SentAt      time.Time
-	CompletedAt time.Time
+	SentAt      int64
+	CompletedAt int64
 	Description string
 	Request     []byte // *sliverpb.Envelope
 	Response    []byte // *sliverpb.Envelope
@@ -152,10 +152,10 @@ func (b *BeaconTask) ToProtobuf(content bool) *clientpb.BeaconTask {
 	task := &clientpb.BeaconTask{
 		ID:          b.ID.String(),
 		BeaconID:    b.BeaconID.String(),
-		CreatedAt:   int64(b.CreatedAt.UTC().Unix()),
+		CreatedAt:   b.CreatedAt.Unix(),
 		State:       b.State,
-		SentAt:      int64(b.SentAt.UTC().Unix()),
-		CompletedAt: int64(b.CompletedAt.UTC().Unix()),
+		SentAt:      b.SentAt,
+		CompletedAt: b.CompletedAt,
 		Description: b.Description,
 	}
 	if content {

server/db/models/implant.go

@@ -66,6 +66,37 @@ func (ib *ImplantBuild) BeforeCreate(tx *gorm.DB) (err error) {
 	return nil
 }
 
+// Convert ImplantBuild To Protobuf
+func (ib *ImplantBuild) ToProtobuf() *clientpb.ImplantBuild {
+	build := clientpb.ImplantBuild{
+		ID:              ib.ID.String(),
+		Name:            ib.Name,
+		MD5:             ib.MD5,
+		SHA1:            ib.SHA1,
+		SHA256:          ib.SHA256,
+		Burned:          ib.Burned,
+		ImplantID:       ib.ImplantID,
+		ImplantConfigID: ib.ImplantConfigID.String(),
+	}
+	return &build
+}
+
+func ImplantBuildFromProtobuf(ib *clientpb.ImplantBuild) *ImplantBuild {
+	id, _ := uuid.FromString(ib.ID)
+	ImplantConfidID, _ := uuid.FromString(ib.ImplantConfigID)
+	build := ImplantBuild{
+		ID:              id,
+		Name:            ib.Name,
+		MD5:             ib.MD5,
+		SHA1:            ib.SHA1,
+		SHA256:          ib.SHA256,
+		Burned:          ib.Burned,
+		ImplantID:       ib.ImplantID,
+		ImplantConfigID: ImplantConfidID,
+	}
+	return &build
+}
+
 // ImplantConfig - An implant build configuration
 type ImplantConfig struct {
 	ID               uuid.UUID `gorm:"primaryKey;->;<-:create;type:uuid;"`

server/generate/implants.go

@@ -110,23 +110,21 @@ func ImplantBuildSave(name string, config *clientpb.ImplantConfig, fPath string)
 		return err
 	}
 
-	id, _ := uuid.FromString(config.ID)
-	dbSession := db.Session()
-	implantBuild := &models.ImplantBuild{
+	implantBuild := &clientpb.ImplantBuild{
 		Name:            name,
-		ImplantConfigID: id,
+		ImplantConfigID: config.ID,
 		MD5:             md5Hash,
 		SHA1:            sha1Hash,
 		SHA256:          sha256Hash,
 		ImplantID:       implantID,
 	}
-	watchtower.AddImplantToWatchlist(implantBuild)
-	result := dbSession.Create(&implantBuild)
-	if result.Error != nil {
-		return result.Error
+	implantBuild, err = db.SaveImplantBuild(implantBuild)
+	if err != nil {
+		return err
 	}
+	watchtower.AddImplantToWatchlist(implantBuild)
 	storageLog.Infof("%s -> %s", implantBuild.ID, implantBuild.Name)
-	return os.WriteFile(filepath.Join(buildsDir, implantBuild.ID.String()), data, 0600)
+	return os.WriteFile(filepath.Join(buildsDir, implantBuild.ID), data, 0600)
 }
 
 func computeHashes(data []byte) (string, string, string) {
@@ -140,25 +138,25 @@ func computeHashes(data []byte) (string, string, string) {
 }
 
 // ImplantFileFromBuild - Saves a binary file into the database
-func ImplantFileFromBuild(build *models.ImplantBuild) ([]byte, error) {
+func ImplantFileFromBuild(build *clientpb.ImplantBuild) ([]byte, error) {
 	buildsDir, err := getBuildsDir()
 	if err != nil {
 		return nil, err
 	}
-	buildFilePath := path.Join(buildsDir, build.ID.String())
+	buildFilePath := path.Join(buildsDir, build.ID)
 	if _, err := os.Stat(buildFilePath); os.IsNotExist(err) {
 		return nil, ErrImplantBuildFileNotFound
 	}
 	return os.ReadFile(buildFilePath)
 }
 
 // ImplantFileDelete - Delete the implant from the file system
-func ImplantFileDelete(build *models.ImplantBuild) error {
+func ImplantFileDelete(build *clientpb.ImplantBuild) error {
 	buildsDir, err := getBuildsDir()
 	if err != nil {
 		return err
 	}
-	buildFilePath := filepath.Join(buildsDir, build.ID.String())
+	buildFilePath := filepath.Join(buildsDir, build.ID)
 	if _, err := os.Stat(buildFilePath); os.IsNotExist(err) {
 		return ErrImplantBuildFileNotFound
 	}

server/handlers/beacons.go

@@ -165,7 +165,7 @@ func beaconTasksHandler(implantConn *core.ImplantConnection, data []byte) *slive
 		envelope.ID = pendingTask.EnvelopeID
 		tasks = append(tasks, envelope)
 		pendingTask.State = models.SENT
-		pendingTask.SentAt = time.Now()
+		pendingTask.SentAt = time.Now().Unix()
 		err = db.Session().Model(&models.BeaconTask{}).Where(&models.BeaconTask{
 			ID: pendingTask.ID,
 		}).Updates(pendingTask).Error

server/rpc/rpc-generate.go

@@ -152,7 +152,7 @@ func (rpc *Server) Regenerate(ctx context.Context, req *clientpb.RegenerateReq)
 	if err != nil {
 		return nil, err
 	}
-	config, err := db.ImplantConfigByID(build.ImplantConfigID.String())
+	config, err := db.ImplantConfigByID(build.ImplantConfigID)
 	if err != nil {
 		return nil, err
 	}

server/rpc/rpc-tasks.go

@@ -310,7 +310,7 @@ func getSliverShellcode(name string) ([]byte, string, error) {
 		return nil, "", err
 	}
 
-	config, err := db.ImplantConfigByID(build.ImplantConfigID.String())
+	config, err := db.ImplantConfigByID(build.ImplantConfigID)
 	if err != nil {
 		return nil, "", err
 	}

server/watchtower/watchtower.go

@@ -19,7 +19,7 @@ var (
 	watchtowerLog = log.NamedLogger("watchtower", "samples")
 )
 
-func update(implantBuild *models.ImplantBuild) {
+func update(implantBuild *clientpb.ImplantBuild) {
 	if watcher != nil && initialized {
 		watchtowerLog.Debugf("Monitoring implant %s (%s)", implantBuild.Name, implantBuild.MD5)
 		watcher.Add(implantBuild.Name, implantBuild.MD5)
@@ -99,7 +99,7 @@ func StartWatchTower(configs *clientpb.MonitoringProviders) error {
 	return nil
 }
 
-func AddImplantToWatchlist(implant *models.ImplantBuild) {
+func AddImplantToWatchlist(implant *clientpb.ImplantBuild) {
 	update(implant)
 }