178d984996
The goroutine running the monitor isn't actually closed. This PR updates the API to allow that Goroutine to properly block program exit. This can be leveraged as we continue to make the configuration non-global.
74 lines
1.7 KiB
Go
74 lines
1.7 KiB
Go
package zgrab2
|
|
|
|
import "sync"
|
|
|
|
// Monitor is a collection of states per scans and a channel to communicate
|
|
// those scans to the monitor
|
|
type Monitor struct {
|
|
states map[string]*State
|
|
statusesChan chan moduleStatus
|
|
// Callback is invoked after each scan.
|
|
Callback func(string)
|
|
}
|
|
|
|
// State contains the respective number of successes and failures
|
|
// for a given scan
|
|
type State struct {
|
|
Successes uint `json:"successes"`
|
|
Failures uint `json:"failures"`
|
|
}
|
|
|
|
type moduleStatus struct {
|
|
name string
|
|
st status
|
|
}
|
|
|
|
type status uint
|
|
|
|
const (
|
|
statusSuccess status = iota
|
|
statusFailure status = iota
|
|
)
|
|
|
|
// GetStatuses returns a mapping from scanner names to the current number
|
|
// of successes and failures for that scanner
|
|
func (m *Monitor) GetStatuses() map[string]*State {
|
|
return m.states
|
|
}
|
|
|
|
// Stop indicates the monitor is done and the internal channel should be closed.
|
|
// This function does not block, but will allow a call to Wait() on the
|
|
// WaitGroup passed to MakeMonitor to return.
|
|
func (m *Monitor) Stop() {
|
|
close(m.statusesChan)
|
|
}
|
|
|
|
// MakeMonitor returns a Monitor object that can be used to collect and send
|
|
// the status of a running scan
|
|
func MakeMonitor(statusChanSize int, wg *sync.WaitGroup) *Monitor {
|
|
m := new(Monitor)
|
|
m.statusesChan = make(chan moduleStatus, statusChanSize)
|
|
m.states = make(map[string]*State, 10)
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
for s := range m.statusesChan {
|
|
if m.states[s.name] == nil {
|
|
m.states[s.name] = new(State)
|
|
}
|
|
if m.Callback != nil {
|
|
m.Callback(s.name)
|
|
}
|
|
switch s.st {
|
|
case statusSuccess:
|
|
m.states[s.name].Successes++
|
|
case statusFailure:
|
|
m.states[s.name].Failures++
|
|
default:
|
|
continue
|
|
}
|
|
}
|
|
}()
|
|
return m
|
|
}
|