Associating IOC Protobuf objects with database objects to fix IOC deletion, changing some error messages

Fix dark mode docs

client/command/hosts/hosts-ioc-rm.go

@@ -26,7 +26,7 @@ import (
 	"github.com/bishopfox/sliver/client/console"
 )
 
-// HostsIOCRmCmd - Remove a host from the database
+// HostsIOCRmCmd - Remove an IOC from the database
 func HostsIOCRmCmd(cmd *cobra.Command, con *console.SliverConsoleClient, args []string) {
 	host, err := SelectHost(con)
 	if err != nil {
@@ -43,5 +43,5 @@ func HostsIOCRmCmd(cmd *cobra.Command, con *console.SliverConsoleClient, args []
 		con.PrintErrorf("%s\n", err)
 		return
 	}
-	con.PrintInfof("Removed host from database\n")
+	con.PrintInfof("Removed IOC from database\n")
 }

client/command/hosts/hosts-ioc.go

@@ -62,6 +62,10 @@ func hostIOCsTable(host *clientpb.Host, con *console.SliverConsoleClient) string
 }
 
 func SelectHostIOC(host *clientpb.Host, con *console.SliverConsoleClient) (*clientpb.IOC, error) {
+	if len(host.IOCs) == 0 {
+		return nil, ErrNoIOCs
+	}
+
 	// Sort the keys because maps have a randomized order, these keys must be ordered for the selection
 	// to work properly since we rely on the index of the user's selection to find the session in the map
 	var keys []string

client/command/hosts/hosts.go

@@ -42,6 +42,8 @@ import (
 var (
 	// ErrNoHosts - No hosts in database
 	ErrNoHosts = errors.New("no hosts")
+	// ErrNoIOCs - No IOCs in database
+	ErrNoIOCs = errors.New("no IOCs in database for selected host")
 	// ErrNoSelection - No selection made
 	ErrNoSelection = errors.New("no selection")
 )

go.mod

@@ -58,7 +58,7 @@ require (
 	gopkg.in/AlecAivazis/survey.v1 v1.8.8
 	gorm.io/driver/mysql v1.5.1
 	gorm.io/driver/postgres v1.5.2
-	gorm.io/driver/sqlite v1.5.2
+	gorm.io/driver/sqlite v1.5.4
 	gorm.io/gorm v1.25.5
 	gvisor.dev/gvisor v0.0.0-20231116203655-c480f66679e6
 	modernc.org/sqlite v1.27.0

go.sum

@@ -620,8 +620,8 @@ gorm.io/driver/mysql v1.5.1 h1:WUEH5VF9obL/lTtzjmML/5e6VfFR/788coz2uaVCAZw=
 gorm.io/driver/mysql v1.5.1/go.mod h1:Jo3Xu7mMhCyj8dlrb3WoCaRd1FhsVh+yMXb1jUInf5o=
 gorm.io/driver/postgres v1.5.2 h1:ytTDxxEv+MplXOfFe3Lzm7SjG09fcdb3Z/c056DTBx0=
 gorm.io/driver/postgres v1.5.2/go.mod h1:fmpX0m2I1PKuR7mKZiEluwrP3hbs+ps7JIGMUBpCgl8=
-gorm.io/driver/sqlite v1.5.2 h1:TpQ+/dqCY4uCigCFyrfnrJnrW9zjpelWVoEVNy5qJkc=
-gorm.io/driver/sqlite v1.5.2/go.mod h1:qxAuCol+2r6PannQDpOP1FP6ag3mKi4esLnB/jHed+4=
+gorm.io/driver/sqlite v1.5.4 h1:IqXwXi8M/ZlPzH/947tn5uik3aYQslP9BVveoax0nV0=
+gorm.io/driver/sqlite v1.5.4/go.mod h1:qxAuCol+2r6PannQDpOP1FP6ag3mKi4esLnB/jHed+4=
 gorm.io/gorm v1.25.1/go.mod h1:L4uxeKpfBml98NYqVqwAdmV1a2nBtAec/cf3fpucW/k=
 gorm.io/gorm v1.25.5 h1:zR9lOiiYf09VNh5Q1gphfyia1JpiClIWG9hQaxB/mls=
 gorm.io/gorm v1.25.5/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8=

server/db/models/host.go

@@ -36,8 +36,8 @@ type Host struct {
 	OSVersion string // Verbose OS version
 	Locale    string // Detected language code
 
-	IOCs          []IOC
-	ExtensionData []ExtensionData
+	IOCs          []IOC           `gorm:"foreignKey:HostID;references:HostUUID"`
+	ExtensionData []ExtensionData `gorm:"foreignKey:HostID;references:HostUUID"`
 }
 
 // BeforeCreate - GORM hook
@@ -107,6 +107,7 @@ func (i *IOC) ToProtobuf() *clientpb.IOC {
 	return &clientpb.IOC{
 		Path:     i.Path,
 		FileHash: i.FileHash,
+		ID:       i.ID.String(),
 	}
 }
 

server/rpc/rpc-filesystem.go

@@ -29,6 +29,7 @@ import (
 	"github.com/bishopfox/sliver/server/db"
 	"github.com/bishopfox/sliver/server/db/models"
 	"github.com/bishopfox/sliver/server/log"
+	"github.com/bishopfox/sliver/util/encoders"
 	"github.com/gofrs/uuid"
 )
 
@@ -189,6 +190,18 @@ func (rpc *Server) MemfilesRm(ctx context.Context, req *sliverpb.MemfilesRmReq)
 	return resp, nil
 }
 
+func hashUploadData(encoder string, data []byte) [32]byte {
+	if encoder == "gzip" {
+		decodedData, err := new(encoders.Gzip).Decode(data)
+		if err != nil {
+			return sha256.Sum256(nil)
+		}
+		return sha256.Sum256(decodedData)
+	} else {
+		return sha256.Sum256(data)
+	}
+}
+
 func trackIOC(req *sliverpb.UploadReq, resp *sliverpb.Upload) {
 	fsLog.Debugf("Adding IOC to database ...")
 	request := req.GetRequest()
@@ -207,8 +220,8 @@ func trackIOC(req *sliverpb.UploadReq, resp *sliverpb.Upload) {
 		return
 	}
 
-	sum := sha256.Sum256(req.Data)
-	id, _ := uuid.FromString(host.ID)
+	sum := hashUploadData(req.Encoder, req.Data)
+	id, _ := uuid.FromString(host.HostUUID)
 	ioc := &models.IOC{
 		HostID:   id,
 		Path:     resp.Path,

server/rpc/rpc.go

@@ -115,10 +115,6 @@ func (rpc *Server) GenericHandler(req GenericRequest, resp GenericResponse) erro
 		return ErrInvalidSessionID
 	}
 
-	// Overwrite unused implant fields before re-serializing
-	request.SessionID = ""
-	request.BeaconID = ""
-
 	reqData, err := proto.Marshal(req)
 	if err != nil {
 		return err

vendor/gorm.io/driver/sqlite/ddlmod.go

@@ -125,7 +125,7 @@ func parseDDL(strs ...string) (*ddl, error) {
 						ColumnTypeValue:   sql.NullString{String: matches[2], Valid: true},
 						PrimaryKeyValue:   sql.NullBool{Valid: true},
 						UniqueValue:       sql.NullBool{Valid: true},
-						NullableValue:     sql.NullBool{Valid: true},
+						NullableValue:     sql.NullBool{Bool: true, Valid: true},
 						DefaultValueValue: sql.NullString{Valid: false},
 					}
 
@@ -175,6 +175,18 @@ func parseDDL(strs ...string) (*ddl, error) {
 	return &result, nil
 }
 
+func (d *ddl) clone() *ddl {
+	copied := new(ddl)
+	*copied = *d
+
+	copied.fields = make([]string, len(d.fields))
+	copy(copied.fields, d.fields)
+	copied.columns = make([]migrator.ColumnType, len(d.columns))
+	copy(copied.columns, d.columns)
+
+	return copied
+}
+
 func (d *ddl) compile() string {
 	if len(d.fields) == 0 {
 		return d.head
@@ -183,6 +195,21 @@ func (d *ddl) compile() string {
 	return fmt.Sprintf("%s (%s)", d.head, strings.Join(d.fields, ","))
 }
 
+func (d *ddl) renameTable(dst, src string) error {
+	tableReg, err := regexp.Compile("\\s*('|`|\")?\\b" + regexp.QuoteMeta(src) + "\\b('|`|\")?\\s*")
+	if err != nil {
+		return err
+	}
+
+	replaced := tableReg.ReplaceAllString(d.head, fmt.Sprintf(" `%s` ", dst))
+	if replaced == d.head {
+		return fmt.Errorf("failed to look up tablename `%s` from DDL head '%s'", src, d.head)
+	}
+
+	d.head = replaced
+	return nil
+}
+
 func (d *ddl) addConstraint(name string, sql string) {
 	reg := regexp.MustCompile("^CONSTRAINT [\"`]?" + regexp.QuoteMeta(name) + "[\"` ]")
 
@@ -240,3 +267,30 @@ func (d *ddl) getColumns() []string {
 	}
 	return res
 }
+
+func (d *ddl) alterColumn(name, sql string) bool {
+	reg := regexp.MustCompile("^(`|'|\"| )" + regexp.QuoteMeta(name) + "(`|'|\"| ) .*?$")
+
+	for i := 0; i < len(d.fields); i++ {
+		if reg.MatchString(d.fields[i]) {
+			d.fields[i] = sql
+			return false
+		}
+	}
+
+	d.fields = append(d.fields, sql)
+	return true
+}
+
+func (d *ddl) removeColumn(name string) bool {
+	reg := regexp.MustCompile("^(`|'|\"| )" + regexp.QuoteMeta(name) + "(`|'|\"| ) .*?$")
+
+	for i := 0; i < len(d.fields); i++ {
+		if reg.MatchString(d.fields[i]) {
+			d.fields = append(d.fields[:i], d.fields[i+1:]...)
+			return true
+		}
+	}
+
+	return false
+}

vendor/gorm.io/driver/sqlite/migrator.go

@@ -3,7 +3,6 @@ package sqlite
 import (
 	"database/sql"
 	"fmt"
-	"regexp"
 	"strings"
 
 	"gorm.io/gorm"
@@ -78,23 +77,16 @@ func (m Migrator) HasColumn(value interface{}, name string) bool {
 
 func (m Migrator) AlterColumn(value interface{}, name string) error {
 	return m.RunWithoutForeignKey(func() error {
-		return m.recreateTable(value, nil, func(rawDDL string, stmt *gorm.Statement) (sql string, sqlArgs []interface{}, err error) {
+		return m.recreateTable(value, nil, func(ddl *ddl, stmt *gorm.Statement) (*ddl, []interface{}, error) {
 			if field := stmt.Schema.LookUpField(name); field != nil {
-				// lookup field from table definition, ddl might looks like `'name' int,` or `'name' int)`
-				reg, err := regexp.Compile("(`|'|\"| )" + field.DBName + "(`|'|\"| ) .*?(,|\\)\\s*$)")
-				if err != nil {
-					return "", nil, err
+				if ddl.alterColumn(field.DBName, fmt.Sprintf("`%s` ?", field.DBName)) {
+					return nil, nil, fmt.Errorf("field `%s` not found in origin ddl, ddl= '%s'", name, ddl.compile())
 				}
 
-				createSQL := reg.ReplaceAllString(rawDDL, fmt.Sprintf("`%v` ?$3", field.DBName))
-
-				if createSQL == rawDDL {
-					return "", nil, fmt.Errorf("failed to look up field %v from DDL %v", field.DBName, rawDDL)
-				}
-
-				return createSQL, []interface{}{m.FullDataTypeOf(field)}, nil
+				return ddl, []interface{}{m.FullDataTypeOf(field)}, nil
 			}
-			return "", nil, fmt.Errorf("failed to alter field with name %v", name)
+
+			return nil, nil, fmt.Errorf("failed to alter field with name `%s`", name)
 		})
 	})
 }
@@ -149,19 +141,13 @@ func (m Migrator) ColumnTypes(value interface{}) ([]gorm.ColumnType, error) {
 }
 
 func (m Migrator) DropColumn(value interface{}, name string) error {
-	return m.recreateTable(value, nil, func(rawDDL string, stmt *gorm.Statement) (sql string, sqlArgs []interface{}, err error) {
+	return m.recreateTable(value, nil, func(ddl *ddl, stmt *gorm.Statement) (*ddl, []interface{}, error) {
 		if field := stmt.Schema.LookUpField(name); field != nil {
 			name = field.DBName
 		}
 
-		reg, err := regexp.Compile("(`|'|\"| |\\[)" + name + "(`|'|\"| |\\]) .*?,")
-		if err != nil {
-			return "", nil, err
-		}
-
-		createSQL := reg.ReplaceAllString(rawDDL, "")
-
-		return createSQL, nil, nil
+		ddl.removeColumn(name)
+		return ddl, nil, nil
 	})
 }
 
@@ -170,7 +156,7 @@ func (m Migrator) CreateConstraint(value interface{}, name string) error {
 		constraint, chk, table := m.GuessConstraintAndTable(stmt, name)
 
 		return m.recreateTable(value, &table,
-			func(rawDDL string, stmt *gorm.Statement) (sql string, sqlArgs []interface{}, err error) {
+			func(ddl *ddl, stmt *gorm.Statement) (*ddl, []interface{}, error) {
 				var (
 					constraintName   string
 					constraintSql    string
@@ -185,17 +171,11 @@ func (m Migrator) CreateConstraint(value interface{}, name string) error {
 					constraintSql = "CONSTRAINT ? CHECK (?)"
 					constraintValues = []interface{}{clause.Column{Name: chk.Name}, clause.Expr{SQL: chk.Constraint}}
 				} else {
-					return "", nil, nil
+					return nil, nil, nil
 				}
 
-				createDDL, err := parseDDL(rawDDL)
-				if err != nil {
-					return "", nil, err
-				}
-				createDDL.addConstraint(constraintName, constraintSql)
-				createSQL := createDDL.compile()
-
-				return createSQL, constraintValues, nil
+				ddl.addConstraint(constraintName, constraintSql)
+				return ddl, constraintValues, nil
 			})
 	})
 }
@@ -210,15 +190,9 @@ func (m Migrator) DropConstraint(value interface{}, name string) error {
 		}
 
 		return m.recreateTable(value, &table,
-			func(rawDDL string, stmt *gorm.Statement) (sql string, sqlArgs []interface{}, err error) {
-				createDDL, err := parseDDL(rawDDL)
-				if err != nil {
-					return "", nil, err
-				}
-				createDDL.removeConstraint(name)
-				createSQL := createDDL.compile()
-
-				return createSQL, nil, nil
+			func(ddl *ddl, stmt *gorm.Statement) (*ddl, []interface{}, error) {
+				ddl.removeConstraint(name)
+				return ddl, nil, nil
 			})
 	})
 }
@@ -375,8 +349,10 @@ func (m Migrator) getRawDDL(table string) (string, error) {
 	return createSQL, nil
 }
 
-func (m Migrator) recreateTable(value interface{}, tablePtr *string,
-	getCreateSQL func(rawDDL string, stmt *gorm.Statement) (sql string, sqlArgs []interface{}, err error)) error {
+func (m Migrator) recreateTable(
+	value interface{}, tablePtr *string,
+	getCreateSQL func(ddl *ddl, stmt *gorm.Statement) (sql *ddl, sqlArgs []interface{}, err error),
+) error {
 	return m.RunWithValue(value, func(stmt *gorm.Statement) error {
 		table := stmt.Table
 		if tablePtr != nil {
@@ -388,27 +364,26 @@ func (m Migrator) recreateTable(value interface{}, tablePtr *string,
 			return err
 		}
 
-		newTableName := table + "__temp"
-
-		createSQL, sqlArgs, err := getCreateSQL(rawDDL, stmt)
+		originDDL, err := parseDDL(rawDDL)
 		if err != nil {
 			return err
 		}
-		if createSQL == "" {
+
+		createDDL, sqlArgs, err := getCreateSQL(originDDL.clone(), stmt)
+		if err != nil {
+			return err
+		}
+		if createDDL == nil {
 			return nil
 		}
 
-		tableReg, err := regexp.Compile("\\s*('|`|\")?\\b" + table + "\\b('|`|\")?\\s*")
-		if err != nil {
+		newTableName := table + "__temp"
+		if err := createDDL.renameTable(newTableName, table); err != nil {
 			return err
 		}
-		createSQL = tableReg.ReplaceAllString(createSQL, fmt.Sprintf(" `%v` ", newTableName))
 
-		createDDL, err := parseDDL(createSQL)
-		if err != nil {
-			return err
-		}
 		columns := createDDL.getColumns()
+		createSQL := createDDL.compile()
 
 		return m.DB.Transaction(func(tx *gorm.DB) error {
 			if err := tx.Exec(createSQL, sqlArgs...).Error; err != nil {

vendor/gorm.io/driver/sqlite/sqlite.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"database/sql"
 	"strconv"
-	"strings"
 
 	"gorm.io/gorm/callbacks"
 
@@ -56,7 +55,7 @@ func (dialector Dialector) Initialize(db *gorm.DB) (err error) {
 	if compareVersion(version, "3.35.0") >= 0 {
 		callbacks.RegisterDefaultCallbacks(db, &callbacks.Config{
 			CreateClauses:        []string{"INSERT", "VALUES", "ON CONFLICT", "RETURNING"},
-			UpdateClauses:        []string{"UPDATE", "SET", "WHERE", "RETURNING"},
+			UpdateClauses:        []string{"UPDATE", "SET", "FROM", "WHERE", "RETURNING"},
 			DeleteClauses:        []string{"DELETE", "FROM", "WHERE", "RETURNING"},
 			LastInsertIDReversed: true,
 		})
@@ -143,19 +142,51 @@ func (dialector Dialector) BindVarTo(writer clause.Writer, stmt *gorm.Statement,
 }
 
 func (dialector Dialector) QuoteTo(writer clause.Writer, str string) {
-	writer.WriteByte('`')
-	if strings.Contains(str, ".") {
-		for idx, str := range strings.Split(str, ".") {
-			if idx > 0 {
-				writer.WriteString(".`")
+	var (
+		underQuoted, selfQuoted bool
+		continuousBacktick      int8
+		shiftDelimiter          int8
+	)
+
+	for _, v := range []byte(str) {
+		switch v {
+		case '`':
+			continuousBacktick++
+			if continuousBacktick == 2 {
+				writer.WriteString("``")
+				continuousBacktick = 0
 			}
-			writer.WriteString(str)
-			writer.WriteByte('`')
+		case '.':
+			if continuousBacktick > 0 || !selfQuoted {
+				shiftDelimiter = 0
+				underQuoted = false
+				continuousBacktick = 0
+				writer.WriteString("`")
+			}
+			writer.WriteByte(v)
+			continue
+		default:
+			if shiftDelimiter-continuousBacktick <= 0 && !underQuoted {
+				writer.WriteString("`")
+				underQuoted = true
+				if selfQuoted = continuousBacktick > 0; selfQuoted {
+					continuousBacktick -= 1
+				}
+			}
+
+			for ; continuousBacktick > 0; continuousBacktick -= 1 {
+				writer.WriteString("``")
+			}
+
+			writer.WriteByte(v)
 		}
-	} else {
-		writer.WriteString(str)
-		writer.WriteByte('`')
+		shiftDelimiter++
 	}
+
+	if continuousBacktick > 0 && !selfQuoted {
+		writer.WriteString("``")
+	}
+	writer.WriteString("`")
 }
 
 func (dialector Dialector) Explain(sql string, vars ...interface{}) string {
@@ -167,7 +198,8 @@ func (dialector Dialector) DataTypeOf(field *schema.Field) string {
 	case schema.Bool:
 		return "numeric"
 	case schema.Int, schema.Uint:
-		if field.AutoIncrement && !field.PrimaryKey {
+		if field.AutoIncrement {
+			// doesn't check `PrimaryKey`, to keep backward compatibility
 			// https://www.sqlite.org/autoinc.html
 			return "integer PRIMARY KEY AUTOINCREMENT"
 		} else {

vendor/modules.txt

@@ -1027,7 +1027,7 @@ gorm.io/driver/mysql
 # gorm.io/driver/postgres v1.5.2
 ## explicit; go 1.18
 gorm.io/driver/postgres
-# gorm.io/driver/sqlite v1.5.2
+# gorm.io/driver/sqlite v1.5.4
 ## explicit; go 1.20
 gorm.io/driver/sqlite
 # gorm.io/gorm v1.25.5