misc: reporting session replay loading progress as api.rest state object

modules/api_rest/api_rest.go

@@ -60,6 +60,7 @@ func NewRestAPI(s *session.Session) *RestAPI {
 	mod.State.Store("recording", &mod.recording)
 	mod.State.Store("replaying", &mod.replaying)
 	mod.State.Store("loading", &mod.loading)
+	mod.State.Store("load_progress", 0)
 	mod.State.Store("rec_time", &mod.recTime)
 	mod.State.Store("rec_filename", &mod.recordFileName)
 	mod.State.Store("rec_frames", 0)

modules/api_rest/api_rest_record.go

@@ -43,7 +43,7 @@ func (mod *RestAPI) recorder() {
 	mod.recTime = 0
 	mod.recording = true
 	mod.replaying = false
-	mod.record = NewRecord(mod.recordFileName)
+	mod.record = NewRecord(mod.recordFileName, &mod.SessionModule)
 
 	mod.Info("started recording to %s ...", mod.recordFileName)
 

modules/api_rest/api_rest_replay.go

@@ -25,6 +25,7 @@ func (mod *RestAPI) startReplay(filename string) (err error) {
 		return err
 	}
 
+	mod.State.Store("load_progress", 0)
 	mod.loading = true
 	defer func() {
 		mod.loading = false
@@ -33,7 +34,7 @@ func (mod *RestAPI) startReplay(filename string) (err error) {
 	mod.Info("loading %s ...", mod.recordFileName)
 
 	start := time.Now()
-	if mod.record, err = LoadRecord(mod.recordFileName); err != nil {
+	if mod.record, err = LoadRecord(mod.recordFileName, &mod.SessionModule); err != nil {
 		return err
 	}
 	loadedIn := time.Since(start)

modules/api_rest/record.go

@@ -9,6 +9,8 @@ import (
 	"os"
 	"sync"
 
+	"github.com/bettercap/bettercap/session"
+
 	"github.com/evilsocket/islazy/fs"
 	"github.com/kr/binarydist"
 )
@@ -16,6 +18,8 @@ import (
 type patch []byte
 type frame []byte
 
+type progressCallback func(done int)
+
 type RecordEntry struct {
 	sync.Mutex
 
@@ -26,9 +30,10 @@ type RecordEntry struct {
 	CurState  int     `json:"-"`
 
 	frames   []frame
+	progress progressCallback
 }
 
-func NewRecordEntry() *RecordEntry {
+func NewRecordEntry(progress progressCallback) *RecordEntry {
 	return &RecordEntry{
 		Data:      nil,
 		Cur:       nil,
@@ -36,6 +41,7 @@ func NewRecordEntry() *RecordEntry {
 		NumStates: 0,
 		CurState:  0,
 		frames:    nil,
+		progress:  progress,
 	}
 }
 
@@ -98,8 +104,12 @@ func (e *RecordEntry) Compile() error {
 
 		e.Cur = newWriter.Bytes()
 		e.frames[i+1] = e.Cur
+
+		e.progress(1)
 	}
 
+	e.progress(1)
+
 	return nil
 }
 
@@ -140,20 +150,34 @@ func (e *RecordEntry) Next() []byte {
 type Record struct {
 	sync.Mutex
 
+	mod      *session.SessionModule `json:"-"`
 	fileName string                 `json:"-"`
+	done     int                    `json:"-"`
+	total    int                    `json:"-"`
+	progress float64                `json:"-"`
 	Session  *RecordEntry           `json:"session"`
 	Events   *RecordEntry           `json:"events"`
 }
 
-func NewRecord(fileName string) *Record {
-	return &Record{
+func NewRecord(fileName string, mod *session.SessionModule) *Record {
+	r := &Record{
 		fileName: fileName,
-		Session:  NewRecordEntry(),
-		Events:   NewRecordEntry(),
+		mod:      mod,
 	}
+
+	r.Session = NewRecordEntry(r.onProgress)
+	r.Events = NewRecordEntry(r.onProgress)
+
+	return r
 }
 
-func LoadRecord(fileName string) (*Record, error) {
+func (r *Record) onProgress(done int) {
+	r.done += done
+	r.progress = float64(r.done) / float64(r.total) * 100.0
+	r.mod.State.Store("load_progress", r.progress)
+}
+
+func LoadRecord(fileName string, mod *session.SessionModule) (*Record, error) {
 	if !fs.Exists(fileName) {
 		return nil, fmt.Errorf("%s does not exist", fileName)
 	}
@@ -182,6 +206,16 @@ func LoadRecord(fileName string) (*Record, error) {
 	}
 
 	rec.fileName = fileName
+	rec.mod = mod
+
+	rec.Session.NumStates = len(rec.Session.States)
+	rec.Session.progress = rec.onProgress
+	rec.Events.NumStates = len(rec.Events.States)
+	rec.Events.progress = rec.onProgress
+
+	rec.done = 0
+	rec.total = rec.Session.NumStates + rec.Events.NumStates + 2
+	rec.progress = 0.0
 
 	// reset state and precompute frames
 	if err = rec.Session.Compile(); err != nil {

session/module.go

@@ -60,7 +60,7 @@ type SessionModule struct {
 	Session    *Session
 	Started    bool
 	StatusLock *sync.RWMutex
-	State      sync.Map
+	State      *sync.Map
 
 	handlers []ModuleHandler
 	params   map[string]*ModuleParam
@@ -77,6 +77,7 @@ func NewSessionModule(name string, s *Session) SessionModule {
 		Session:    s,
 		Started:    false,
 		StatusLock: &sync.RWMutex{},
+		State:      &sync.Map{},
 
 		handlers: make([]ModuleHandler, 0),
 		params:   make(map[string]*ModuleParam),