diff --git a/server/db/helpers.go b/server/db/helpers.go index 196f50676a..5d7a9c54b9 100644 --- a/server/db/helpers.go +++ b/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 diff --git a/server/db/models/host.go b/server/db/models/host.go index d16ca4396e..578aca7d35 100644 --- a/server/db/models/host.go +++ b/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 diff --git a/server/db/models/http-c2.go b/server/db/models/http-c2.go index 18d816b036..d45d72ac7c 100644 --- a/server/db/models/http-c2.go +++ b/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 diff --git a/server/db/models/implant.go b/server/db/models/implant.go index cccda8d286..2c244cab6d 100644 --- a/server/db/models/implant.go +++ b/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 diff --git a/server/db/sql.go b/server/db/sql.go index 673e53ab4e..6cb730f47f 100644 --- a/server/db/sql.go +++ b/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 diff --git a/server/generate/implants.go b/server/generate/implants.go index 13169904fc..2d9e6ef34a 100644 --- a/server/generate/implants.go +++ b/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 diff --git a/server/generate/profiles.go b/server/generate/profiles.go index eed64c18c9..f8e8417f1d 100644 --- a/server/generate/profiles.go +++ b/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)