From 695657fbadc516c926c7c70e912f6d1b3478aeca Mon Sep 17 00:00:00 2001 From: Matthijs Gielen Date: Tue, 18 Jun 2024 11:13:49 +0200 Subject: [PATCH] Implement saving and retrieval of session tasks --- client/command/tasks/commands.go | 3 +- client/command/tasks/fetch.go | 2 +- client/command/tasks/helpers.go | 16 +++++-- client/command/tasks/tasks.go | 2 +- client/console/console.go | 12 +++++ server/rpc/rpc-beacons.go | 6 +-- server/rpc/rpc.go | 80 ++++++++++++++++++++------------ 7 files changed, 79 insertions(+), 42 deletions(-) diff --git a/client/command/tasks/commands.go b/client/command/tasks/commands.go index d92e1b2df0..3a6e2c8116 100644 --- a/client/command/tasks/commands.go +++ b/client/command/tasks/commands.go @@ -19,8 +19,7 @@ func Commands(con *console.SliverClient) []*cobra.Command { Run: func(cmd *cobra.Command, args []string) { TasksCmd(cmd, con, args) }, - GroupID: consts.SliverCoreHelpGroup, - Annotations: flags.RestrictTargets(consts.BeaconCmdsFilter), + GroupID: consts.SliverCoreHelpGroup, } flags.Bind("tasks", true, tasksCmd, func(f *pflag.FlagSet) { f.IntP("timeout", "t", flags.DefaultTimeout, "grpc timeout in seconds") diff --git a/client/command/tasks/fetch.go b/client/command/tasks/fetch.go index 3aa8880550..4a88d4c748 100644 --- a/client/command/tasks/fetch.go +++ b/client/command/tasks/fetch.go @@ -47,7 +47,7 @@ import ( // TasksFetchCmd - Manage beacon tasks. func TasksFetchCmd(cmd *cobra.Command, con *console.SliverClient, args []string) { - beacon := con.ActiveTarget.GetBeaconInteractive() + beacon := con.ActiveTarget.GetSessionOrBeaconInteractive() if beacon == nil { return } diff --git a/client/command/tasks/helpers.go b/client/command/tasks/helpers.go index 6cdf28cd49..549861ee5d 100644 --- a/client/command/tasks/helpers.go +++ b/client/command/tasks/helpers.go @@ -51,12 +51,22 @@ func SelectBeaconTask(tasks []*clientpb.BeaconTask) (*clientpb.BeaconTask, error // BeaconTaskIDCompleter returns a structured list of tasks completions, grouped by state. func BeaconTaskIDCompleter(con *console.SliverClient) carapace.Action { callback := func(ctx carapace.Context) carapace.Action { + id := "" beacon := con.ActiveTarget.GetBeacon() - if beacon == nil { - return carapace.ActionMessage("no active beacon") + if beacon != nil { + id = beacon.ID + } else { + session := con.ActiveTarget.GetSession() + if session != nil { + id = session.ID + } } - beaconTasks, err := con.Rpc.GetBeaconTasks(context.Background(), &clientpb.Beacon{ID: beacon.ID}) + if id == "" { + return carapace.ActionMessage("no active beacon or session") + } + + beaconTasks, err := con.Rpc.GetBeaconTasks(context.Background(), &clientpb.Beacon{ID: id}) if err != nil { return carapace.ActionMessage("Failed to fetch tasks: %s", err.Error()) } diff --git a/client/command/tasks/tasks.go b/client/command/tasks/tasks.go index bd025867e6..6112df4060 100644 --- a/client/command/tasks/tasks.go +++ b/client/command/tasks/tasks.go @@ -33,7 +33,7 @@ import ( // TasksCmd - Manage beacon tasks. func TasksCmd(cmd *cobra.Command, con *console.SliverClient, args []string) { - beacon := con.ActiveTarget.GetBeaconInteractive() + beacon := con.ActiveTarget.GetSessionOrBeaconInteractive() if beacon == nil { return } diff --git a/client/console/console.go b/client/console/console.go index b56f719d81..8b4d2e74f8 100644 --- a/client/console/console.go +++ b/client/console/console.go @@ -715,6 +715,18 @@ func (s *ActiveTarget) GetBeaconInteractive() *clientpb.Beacon { return s.beacon } +// GetSessionOrBeaconInteractive +func (s *ActiveTarget) GetSessionOrBeaconInteractive() *clientpb.Beacon { + if s.beacon != nil { + return s.beacon + } else if s.session != nil { + return &clientpb.Beacon{ID: s.session.ID} + } else { + fmt.Printf(Warn + "Please select a beacon or session via `use`\n") + return nil + } +} + // GetBeacon - Same as GetBeacon() but doesn't print a warning. func (s *ActiveTarget) GetBeacon() *clientpb.Beacon { return s.beacon diff --git a/server/rpc/rpc-beacons.go b/server/rpc/rpc-beacons.go index ad1208134d..ec1813998b 100644 --- a/server/rpc/rpc-beacons.go +++ b/server/rpc/rpc-beacons.go @@ -85,11 +85,7 @@ func (rpc *Server) RmBeacon(ctx context.Context, req *clientpb.Beacon) (*commonp // GetBeaconTasks - Get a list of tasks for a specific beacon func (rpc *Server) GetBeaconTasks(ctx context.Context, req *clientpb.Beacon) (*clientpb.BeaconTasks, error) { - beacon, err := db.BeaconByID(req.ID) - if err != nil { - return nil, ErrInvalidBeaconID - } - tasks, err := db.BeaconTasksByBeaconID(beacon.ID.String()) + tasks, err := db.BeaconTasksByBeaconID(req.ID) return &clientpb.BeaconTasks{Tasks: tasks}, err } diff --git a/server/rpc/rpc.go b/server/rpc/rpc.go index 2f225cbf7e..307ad5707f 100644 --- a/server/rpc/rpc.go +++ b/server/rpc/rpc.go @@ -25,6 +25,7 @@ import ( "strings" "time" + consts "github.com/bishopfox/sliver/client/constants" "github.com/bishopfox/sliver/client/version" "github.com/bishopfox/sliver/protobuf/clientpb" "github.com/bishopfox/sliver/protobuf/commonpb" @@ -32,7 +33,9 @@ import ( "github.com/bishopfox/sliver/protobuf/sliverpb" "github.com/bishopfox/sliver/server/core" "github.com/bishopfox/sliver/server/db" + "github.com/bishopfox/sliver/server/db/models" "github.com/bishopfox/sliver/server/log" + "github.com/gofrs/uuid" "google.golang.org/grpc/credentials" "google.golang.org/grpc/peer" "google.golang.org/protobuf/proto" @@ -104,22 +107,49 @@ func (rpc *Server) GenericHandler(req GenericRequest, resp GenericResponse) erro if request == nil { return ErrMissingRequestField } + reqData, err := proto.Marshal(req) + if err != nil { + return err + } + + taskResponse := resp.GetResponse() + taskResponse.Async = request.Async + beacon := models.Beacon{} + if request.BeaconID != "" { + beacon.ID, err = uuid.FromString(request.BeaconID) + } else if request.SessionID != "" { + beacon.ID, err = uuid.FromString(request.SessionID) + } + task, err := beacon.Task(&sliverpb.Envelope{ + Type: sliverpb.MsgNumber(req), + Data: reqData, + }) + if err != nil { + rpcLog.Errorf("Database error: %s", err) + return ErrDatabaseFailure + } + parts := strings.Split(string(req.ProtoReflect().Descriptor().FullName().Name()), ".") + name := parts[len(parts)-1] + task.Description = name + err = db.Session().Save(task).Error + if err != nil { + rpcLog.Errorf("Database error: %s", err) + return ErrDatabaseFailure + } + rpcLog.Warningf("Task: %#v", task) + if request.Async { - err = rpc.asyncGenericHandler(req, resp) + err = rpc.asyncGenericHandler(req, resp, task) return err } + task.SentAt = time.Now().Unix() // Sync request session := core.Sessions.Get(request.SessionID) if session == nil { return ErrInvalidSessionID } - reqData, err := proto.Marshal(req) - if err != nil { - return err - } - data, err := session.Request(sliverpb.MsgNumber(req), rpc.getTimeout(req), reqData) if err != nil { return err @@ -128,11 +158,24 @@ func (rpc *Server) GenericHandler(req GenericRequest, resp GenericResponse) erro if err != nil { return err } + task.State = models.COMPLETED + task.CompletedAt = time.Now().Unix() + task.Response = data + + err = db.Session().Updates(task).Error + if err != nil { + rpcLog.Errorf("Error updating db task: %s", err) + } + eventData, _ := proto.Marshal(task.ToProtobuf(false)) + core.EventBroker.Publish(core.Event{ + EventType: consts.BeaconTaskResultEvent, + Data: eventData, + }) return rpc.getError(resp) } // asyncGenericHandler - Generic handler for async request/response's for beacon tasks -func (rpc *Server) asyncGenericHandler(req GenericRequest, resp GenericResponse) error { +func (rpc *Server) asyncGenericHandler(req GenericRequest, resp GenericResponse, task *models.BeaconTask) error { // VERY VERBOSE // rpcLog.Debugf("Async Generic Handler: %#v", req) request := req.GetRequest() @@ -146,32 +189,9 @@ func (rpc *Server) asyncGenericHandler(req GenericRequest, resp GenericResponse) return ErrInvalidBeaconID } - // Overwrite unused implant fields before re-serializing - request.SessionID = "" - request.BeaconID = "" - reqData, err := proto.Marshal(req) - if err != nil { - return err - } taskResponse := resp.GetResponse() taskResponse.Async = true taskResponse.BeaconID = beacon.ID.String() - task, err := beacon.Task(&sliverpb.Envelope{ - Type: sliverpb.MsgNumber(req), - Data: reqData, - }) - if err != nil { - rpcLog.Errorf("Database error: %s", err) - return ErrDatabaseFailure - } - parts := strings.Split(string(req.ProtoReflect().Descriptor().FullName().Name()), ".") - name := parts[len(parts)-1] - task.Description = name - err = db.Session().Save(task).Error - if err != nil { - rpcLog.Errorf("Database error: %s", err) - return ErrDatabaseFailure - } taskResponse.TaskID = task.ID.String() rpcLog.Debugf("Successfully tasked beacon: %#v", taskResponse) return nil