diff --git a/modules/api_rest/api_rest.go b/modules/api_rest/api_rest.go index 23b259e0..b20e8030 100644 --- a/modules/api_rest/api_rest.go +++ b/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) diff --git a/modules/api_rest/api_rest_record.go b/modules/api_rest/api_rest_record.go index ab0be552..22f03cd6 100644 --- a/modules/api_rest/api_rest_record.go +++ b/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) diff --git a/modules/api_rest/api_rest_replay.go b/modules/api_rest/api_rest_replay.go index 035b1d3c..b7072d3e 100644 --- a/modules/api_rest/api_rest_replay.go +++ b/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) diff --git a/modules/api_rest/record.go b/modules/api_rest/record.go index 1cc734a4..4693ce2a 100644 --- a/modules/api_rest/record.go +++ b/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 @@ -25,10 +29,11 @@ type RecordEntry struct { NumStates int `json:"-"` CurState int `json:"-"` - frames []frame + 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 - fileName string `json:"-"` - Session *RecordEntry `json:"session"` - Events *RecordEntry `json:"events"` + 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 { diff --git a/session/module.go b/session/module.go index 149552d1..f8ba5e75 100644 --- a/session/module.go +++ b/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),