nuclei/pkg/utils/stats/stats.go
mzack9999 06b0219839 lint
2024-05-16 13:48:28 +02:00

130 lines
2.8 KiB
Go

package stats
import (
"fmt"
"sync/atomic"
"github.com/logrusorgru/aurora"
"github.com/projectdiscovery/gologger"
mapsutil "github.com/projectdiscovery/utils/maps"
)
// Storage is a storage for storing statistics information
// about the nuclei engine displaying it at user-defined intervals.
type Storage struct {
data *mapsutil.SyncLockMap[string, *storageDataItem]
}
type storageDataItem struct {
description string
value atomic.Int64
}
var Default *Storage
func init() {
Default = New()
}
// NewEntry creates a new entry in the storage object
func NewEntry(name, description string) {
Default.NewEntry(name, description)
}
// Increment increments the value for a name string
func Increment(name string) {
Default.Increment(name)
}
// Display displays the stats for a name
func Display(name string) {
Default.Display(name)
}
func DisplayAsWarning(name string) {
Default.DisplayAsWarning(name)
}
// ForceDisplayWarning forces the display of a warning
// regardless of current verbosity level
func ForceDisplayWarning(name string) {
Default.ForceDisplayWarning(name)
}
// GetValue returns the value for a set variable
func GetValue(name string) int64 {
return Default.GetValue(name)
}
// New creates a new storage object
func New() *Storage {
data := mapsutil.NewSyncLockMap[string, *storageDataItem]()
return &Storage{data: data}
}
// NewEntry creates a new entry in the storage object
func (s *Storage) NewEntry(name, description string) {
_ = s.data.Set(name, &storageDataItem{description: description, value: atomic.Int64{}})
}
// Increment increments the value for a name string
func (s *Storage) Increment(name string) {
data, ok := s.data.Get(name)
if !ok {
return
}
data.value.Add(1)
}
// Display displays the stats for a name
func (s *Storage) Display(name string) {
data, ok := s.data.Get(name)
if !ok {
return
}
dataValue := data.value.Load()
if dataValue == 0 {
return // don't show for nil stats
}
gologger.Error().Label("WRN").Msgf(data.description, dataValue)
}
func (s *Storage) DisplayAsWarning(name string) {
data, ok := s.data.Get(name)
if !ok {
return
}
dataValue := data.value.Load()
if dataValue == 0 {
return // don't show for nil stats
}
gologger.Warning().Label("WRN").Msgf(data.description, dataValue)
}
// ForceDisplayWarning forces the display of a warning
// regardless of current verbosity level
func (s *Storage) ForceDisplayWarning(name string) {
data, ok := s.data.Get(name)
if !ok {
return
}
dataValue := data.value.Load()
if dataValue == 0 {
return // don't show for nil stats
}
gologger.Print().Msgf("[%v] %v", aurora.BrightYellow("WRN"), fmt.Sprintf(data.description, dataValue))
}
// GetValue returns the value for a set variable
func (s *Storage) GetValue(name string) int64 {
data, ok := s.data.Get(name)
if !ok {
return 0
}
return data.value.Load()
}