Improving PostgreSQL Support

server/db/helpers.go

@@ -27,6 +27,7 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	"errors"
+	"fmt"
 	"os"
 	"path/filepath"
 	"strings"
@@ -258,7 +259,7 @@ func ImplantProfileByName(name string) (*clientpb.ImplantProfile, error) {
 		return nil, err
 	}
 	err = Session().Where(models.ImplantConfig{
-		ImplantProfileID: profile.ID,
+		ImplantProfileID: &profile.ID,
 	}).First(&config).Error
 	if err != nil {
 		return nil, err
@@ -367,7 +368,7 @@ func LoadHTTPC2ConfigByName(name string) (*clientpb.HTTPC2Config, error) {
 	// load headers
 	c2ImplantHeaders := []models.HttpC2Header{}
 	err = Session().Where(&models.HttpC2Header{
-		HttpC2ImplantConfigID: c2ImplantConfig.ID,
+		HttpC2ImplantConfigID: &c2ImplantConfig.ID,
 	}).Find(&c2ImplantHeaders).Error
 	if err != nil {
 		return nil, err
@@ -397,7 +398,7 @@ func LoadHTTPC2ConfigByName(name string) (*clientpb.HTTPC2Config, error) {
 	// load headers
 	c2ServerHeaders := []models.HttpC2Header{}
 	err = Session().Where(&models.HttpC2Header{
-		HttpC2ServerConfigID: c2ServerConfig.ID,
+		HttpC2ServerConfigID: &c2ServerConfig.ID,
 	}).Find(&c2ServerHeaders).Error
 	if err != nil {
 		return nil, err
@@ -471,7 +472,7 @@ func HTTPC2ConfigUpdate(newConf *clientpb.HTTPC2Config, oldConf *clientpb.HTTPC2
 	}
 
 	err = Session().Where(&models.HttpC2Header{
-		HttpC2ServerConfigID: serverID,
+		HttpC2ServerConfigID: &serverID,
 	}).Delete(&models.HttpC2Header{})
 	if err.Error != nil {
 		return err.Error
@@ -495,7 +496,7 @@ func HTTPC2ConfigUpdate(newConf *clientpb.HTTPC2Config, oldConf *clientpb.HTTPC2
 	}
 
 	for _, header := range c2Config.ServerConfig.Headers {
-		header.HttpC2ServerConfigID = serverID
+		header.HttpC2ServerConfigID = &serverID
 		err = Session().Clauses(clause.OnConflict{
 			UpdateAll: true,
 		}).Create(&header)
@@ -591,6 +592,35 @@ func ListenerJobs() ([]*clientpb.ListenerJob, error) {
 }
 
 func DeleteListener(JobID uint32) error {
+	// Determine which type of listener this is and delete its record from the corresponding table
+	var deleteErr error
+	// JobID is unique so if the JobID exists, it is the only record in the table with that JobID
+	listener := &models.ListenerJob{}
+	result := Session().Where(&models.ListenerJob{JobID: JobID}).First(&listener)
+	if result.Error != nil {
+		return result.Error
+	}
+	if listener == nil {
+		return fmt.Errorf("Job ID %d not found in database", JobID)
+	}
+	listenerID := listener.ID
+	switch listener.Type {
+	case constants.HttpStr, constants.HttpsStr:
+		deleteErr = Session().Where(&models.HTTPListener{ListenerJobID: listenerID}).Delete(&models.HTTPListener{}).Error
+	case constants.DnsStr:
+		deleteErr = Session().Where(&models.DNSListener{ListenerJobID: listenerID}).Delete(&models.DNSListener{}).Error
+	case constants.MtlsStr:
+		deleteErr = Session().Where(&models.MtlsListener{ListenerJobID: listenerID}).Delete(&models.MtlsListener{}).Error
+	case constants.WGStr:
+		deleteErr = Session().Where(&models.WGListener{ListenerJobID: listenerID}).Delete(&models.WGListener{}).Error
+	case constants.MultiplayerModeStr:
+		deleteErr = Session().Where(&models.MultiplayerListener{ListenerJobID: listenerID}).Delete(&models.MultiplayerListener{}).Error
+	}
+
+	if deleteErr != nil {
+		return deleteErr
+	}
+
 	return Session().Where(&models.ListenerJob{JobID: JobID}).Delete(&models.ListenerJob{}).Error
 }
 
@@ -641,6 +671,9 @@ func DeleteProfile(name string) error {
 
 	uuid, _ := uuid.FromString(profile.Config.ID)
 
+	// delete linked ImplantC2
+	err = Session().Where(&models.ImplantC2{ImplantConfigID: uuid}).Delete(&models.ImplantC2{}).Error
+
 	// delete linked ImplantConfig
 	err = Session().Where(&models.ImplantConfig{ID: uuid}).Delete(&models.ImplantConfig{}).Error
 

server/db/models/host.go

@@ -29,7 +29,7 @@ import (
 // Host - Represents a host machine
 type Host struct {
 	ID        uuid.UUID `gorm:"primaryKey;->;<-:create;type:uuid;"`
-	HostUUID  uuid.UUID `gorm:"type:uuid;"`
+	HostUUID  uuid.UUID `gorm:"type:uuid;unique"`
 	CreatedAt time.Time `gorm:"->;<-:create;"`
 
 	Hostname  string

server/db/models/http-c2.go

@@ -203,9 +203,9 @@ func (h *HttpC2Cookie) ToProtobuf() *clientpb.HTTPC2Cookie {
 
 // HttpC2Header - HTTP C2 Header (server and implant)
 type HttpC2Header struct {
-	ID                    uuid.UUID `gorm:"primaryKey;->;<-:create;type:uuid;"`
-	HttpC2ServerConfigID  uuid.UUID `gorm:"type:uuid;"`
-	HttpC2ImplantConfigID uuid.UUID `gorm:"type:uuid;"`
+	ID                    uuid.UUID  `gorm:"primaryKey;->;<-:create;type:uuid;"`
+	HttpC2ServerConfigID  *uuid.UUID `gorm:"type:uuid;"`
+	HttpC2ImplantConfigID *uuid.UUID `gorm:"type:uuid;"`
 
 	Method      string
 	Name        string

server/db/models/implant.go

@@ -144,7 +144,7 @@ func ImplantBuildFromProtobuf(ib *clientpb.ImplantBuild) *ImplantBuild {
 // ImplantConfig - An implant build configuration
 type ImplantConfig struct {
 	ID               uuid.UUID `gorm:"primaryKey;->;<-:create;type:uuid;"`
-	ImplantProfileID uuid.UUID
+	ImplantProfileID *uuid.UUID
 
 	ImplantBuilds []ImplantBuild
 	CreatedAt     time.Time `gorm:"->;<-:create;"`
@@ -237,9 +237,8 @@ func (ic *ImplantConfig) ToProtobuf() *clientpb.ImplantConfig {
 		implantBuilds = append(implantBuilds, implantBuild.ToProtobuf())
 	}
 	config := &clientpb.ImplantConfig{
-		ID:               ic.ID.String(),
-		ImplantBuilds:    implantBuilds,
-		ImplantProfileID: ic.ImplantProfileID.String(),
+		ID:            ic.ID.String(),
+		ImplantBuilds: implantBuilds,
 
 		IsBeacon:       ic.IsBeacon,
 		BeaconInterval: ic.BeaconInterval,
@@ -286,6 +285,11 @@ func (ic *ImplantConfig) ToProtobuf() *clientpb.ImplantConfig {
 		IncludeWG:       ic.IncludeWG,
 		IncludeTCP:      ic.IncludeTCP,
 	}
+
+	if ic.ImplantProfileID != nil {
+		config.ImplantProfileID = ic.ImplantProfileID.String()
+	}
+
 	// Copy Canary Domains
 	config.CanaryDomains = []string{}
 	for _, canaryDomain := range ic.CanaryDomains {
@@ -432,7 +436,11 @@ func ImplantConfigFromProtobuf(pbConfig *clientpb.ImplantConfig) *ImplantConfig
 	cfg.ID = id
 	cfg.ImplantBuilds = implantBuilds
 	profileID, _ := uuid.FromString(pbConfig.ImplantProfileID)
-	cfg.ImplantProfileID = profileID
+	if profileID == uuid.Nil {
+		cfg.ImplantProfileID = nil
+	} else {
+		cfg.ImplantProfileID = &profileID
+	}
 
 	cfg.IsBeacon = pbConfig.IsBeacon
 	cfg.BeaconInterval = pbConfig.BeaconInterval

server/db/sql.go

@@ -50,7 +50,14 @@ func newDBClient() *gorm.DB {
 		panic(fmt.Sprintf("Unknown DB Dialect: '%s'", dbConfig.Dialect))
 	}
 
-	err := dbClient.AutoMigrate(
+	// We cannot pass all of these into AutoMigrate at once because if one fails, subsequent models will not be created
+	var allDBModels []interface{} = append(make([]interface{}, 0),
+		&models.HttpC2Header{},
+		&models.HttpC2ServerConfig{},
+		&models.HttpC2ImplantConfig{},
+		&models.HttpC2Config{},
+		&models.HttpC2URLParameter{},
+		&models.HttpC2PathSegment{},
 		&models.Beacon{},
 		&models.BeaconTask{},
 		&models.DNSCanary{},
@@ -62,30 +69,25 @@ func newDBClient() *gorm.DB {
 		&models.CrackFileChunk{},
 		&models.Certificate{},
 		&models.Host{},
-		&models.HttpC2Config{},
-		&models.HttpC2ImplantConfig{},
-		&models.HttpC2ServerConfig{},
-		&models.HttpC2Header{},
-		&models.HttpC2PathSegment{},
+		&models.KeyValue{},
+		&models.WGKeys{},
+		&models.WGPeer{},
+		&models.ResourceID{},
 		&models.HttpC2Cookie{},
-		&models.HttpC2URLParameter{},
 		&models.IOC{},
 		&models.ExtensionData{},
-		&models.ImplantBuild{},
 		&models.ImplantProfile{},
 		&models.ImplantConfig{},
+		&models.ImplantBuild{},
 		&models.ImplantC2{},
 		&models.EncoderAsset{},
 		&models.KeyExHistory{},
-		&models.KeyValue{},
 		&models.CanaryDomain{},
 		&models.Loot{},
 		&models.Credential{},
 		&models.Operator{},
 		&models.Website{},
 		&models.WebContent{},
-		&models.WGKeys{},
-		&models.WGPeer{},
 		&models.ListenerJob{},
 		&models.HTTPListener{},
 		&models.DNSListener{},
@@ -94,10 +96,15 @@ func newDBClient() *gorm.DB {
 		&models.MtlsListener{},
 		&models.DnsDomain{},
 		&models.MonitoringProvider{},
-		&models.ResourceID{},
 	)
-	if err != nil {
-		clientLog.Error(err)
+
+	var err error
+
+	for _, model := range allDBModels {
+		err = dbClient.AutoMigrate(model)
+		if err != nil {
+			clientLog.Error(err)
+		}
 	}
 
 	// Get generic database object sql.DB to use its functions

server/generate/implants.go

@@ -77,7 +77,11 @@ func ImplantConfigSave(config *clientpb.ImplantConfig) (*clientpb.ImplantConfig,
 
 	} else {
 		id, _ := uuid.FromString(dbConfig.ImplantProfileID)
-		modelConfig.ImplantProfileID = id
+		if id == uuid.Nil {
+			modelConfig.ImplantProfileID = nil
+		} else {
+			modelConfig.ImplantProfileID = &id
+		}
 
 		// this avoids gorm saving duplicate c2 objects ...
 		tempC2 := modelConfig.C2

server/generate/profiles.go

@@ -54,7 +54,11 @@ func SaveImplantProfile(pbProfile *clientpb.ImplantProfile) (*clientpb.ImplantPr
 		profile.ImplantConfig.ID = configID
 
 		profileID, _ := uuid.FromString(dbProfile.ID)
-		profile.ImplantConfig.ImplantProfileID = profileID
+		if profileID == uuid.Nil {
+			profile.ImplantConfig.ImplantProfileID = nil
+		} else {
+			profile.ImplantConfig.ImplantProfileID = &profileID
+		}
 
 		for _, c2 := range dbProfile.Config.C2 {
 			id, _ := uuid.FromString(c2.ID)