mirror of
https://github.com/safedep/pmg.git
synced 2026-08-03 07:24:09 +02:00
feat: add progress bar and add debug & silent flags
This commit is contained in:
+17
-1
@@ -8,6 +8,7 @@ import (
|
|||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/safedep/dry/log"
|
"github.com/safedep/dry/log"
|
||||||
|
"github.com/safedep/pmg/internal/ui"
|
||||||
"github.com/safedep/pmg/pkg/analyser"
|
"github.com/safedep/pmg/pkg/analyser"
|
||||||
"github.com/safedep/pmg/pkg/common/utils"
|
"github.com/safedep/pmg/pkg/common/utils"
|
||||||
"github.com/safedep/pmg/pkg/models"
|
"github.com/safedep/pmg/pkg/models"
|
||||||
@@ -19,6 +20,7 @@ import (
|
|||||||
var (
|
var (
|
||||||
packageName string
|
packageName string
|
||||||
action string
|
action string
|
||||||
|
silentScan bool
|
||||||
)
|
)
|
||||||
|
|
||||||
func NewNpmCommand() *cobra.Command {
|
func NewNpmCommand() *cobra.Command {
|
||||||
@@ -49,10 +51,18 @@ func NewNpmCommand() *cobra.Command {
|
|||||||
return utils.ExecCmd(npmPath, args, []string{})
|
return utils.ExecCmd(npmPath, args, []string{})
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
cmd.Flags().BoolVarP(&silentScan, "silent", "s", false,
|
||||||
|
"Silent scan to prevent rendering UI")
|
||||||
return cmd
|
return cmd
|
||||||
}
|
}
|
||||||
|
|
||||||
func wrapNpm() error {
|
func wrapNpm() error {
|
||||||
|
if !silentScan {
|
||||||
|
ui.StartProgressWriter()
|
||||||
|
}
|
||||||
|
var progressTracker ui.ProgressTracker
|
||||||
|
|
||||||
|
progressTracker = ui.TrackProgress(fmt.Sprintf("Scanning %s ", packageName), 1)
|
||||||
if packageName == "" {
|
if packageName == "" {
|
||||||
return fmt.Errorf("package name cannot be empty")
|
return fmt.Errorf("package name cannot be empty")
|
||||||
}
|
}
|
||||||
@@ -81,17 +91,21 @@ func wrapNpm() error {
|
|||||||
// Update packageName with resolved version for npm installation
|
// Update packageName with resolved version for npm installation
|
||||||
packageName = fmt.Sprintf("%s@%s", name, version)
|
packageName = fmt.Sprintf("%s@%s", name, version)
|
||||||
}
|
}
|
||||||
|
ui.IncrementProgress(progressTracker, 1)
|
||||||
|
|
||||||
deps, err := npmFetcher.GetFlattenedDependencies(ctx, name, version)
|
deps, err := npmFetcher.GetFlattenedDependencies(ctx, name, version)
|
||||||
|
ui.IncrementProgress(progressTracker, 1)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
ui.IncrementTrackerTotal(progressTracker, int64(len(deps)))
|
||||||
client, err := analyser.GetMalwareAnalysisClient()
|
client, err := analyser.GetMalwareAnalysisClient()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("error while creating a malware analysis client: %w", err)
|
return fmt.Errorf("error while creating a malware analysis client: %w", err)
|
||||||
}
|
}
|
||||||
pkgAnalyser := analyser.New(client, ctx)
|
pkgAnalyser := analyser.New(client, ctx)
|
||||||
|
|
||||||
|
pkgAnalyser.ProgressTracker = progressTracker
|
||||||
handler := pkgAnalyser.Handler()
|
handler := pkgAnalyser.Handler()
|
||||||
|
|
||||||
// Create work queue with appropriate buffer size and concurrency
|
// Create work queue with appropriate buffer size and concurrency
|
||||||
@@ -114,6 +128,8 @@ func wrapNpm() error {
|
|||||||
|
|
||||||
// Wait for all analysis to complete
|
// Wait for all analysis to complete
|
||||||
queue.Wait()
|
queue.Wait()
|
||||||
|
ui.MarkTrackerAsDone(progressTracker)
|
||||||
|
ui.StopProgressWriter()
|
||||||
|
|
||||||
// Get the npm PATH and continue with installation
|
// Get the npm PATH and continue with installation
|
||||||
npmPath, err := utils.GetExecutablePath("npm")
|
npmPath, err := utils.GetExecutablePath("npm")
|
||||||
|
|||||||
@@ -0,0 +1,136 @@
|
|||||||
|
package ui
|
||||||
|
|
||||||
|
import (
|
||||||
|
"os"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/jedib0t/go-pretty/v6/progress"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ProgressTracker defines the interface for tracking progress
|
||||||
|
type ProgressTracker interface {
|
||||||
|
Increment(count int64)
|
||||||
|
SetValue(count int64)
|
||||||
|
UpdateTotal(count int64)
|
||||||
|
MarkAsDone()
|
||||||
|
GetValue() int64
|
||||||
|
GetTotal() int64
|
||||||
|
}
|
||||||
|
|
||||||
|
// progressTrackerImpl implements ProgressTracker using go-pretty/progress
|
||||||
|
type progressTrackerImpl struct {
|
||||||
|
tracker *progress.Tracker
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *progressTrackerImpl) Increment(count int64) {
|
||||||
|
if p.tracker != nil {
|
||||||
|
p.tracker.Increment(count)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *progressTrackerImpl) SetValue(count int64) {
|
||||||
|
if p.tracker != nil {
|
||||||
|
p.tracker.SetValue(count)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *progressTrackerImpl) UpdateTotal(count int64) {
|
||||||
|
if p.tracker != nil {
|
||||||
|
p.tracker.UpdateTotal(count)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *progressTrackerImpl) MarkAsDone() {
|
||||||
|
if p.tracker != nil {
|
||||||
|
p.tracker.MarkAsDone()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *progressTrackerImpl) GetValue() int64 {
|
||||||
|
if p.tracker != nil {
|
||||||
|
return p.tracker.Value()
|
||||||
|
}
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *progressTrackerImpl) GetTotal() int64 {
|
||||||
|
if p.tracker != nil {
|
||||||
|
return p.tracker.Total
|
||||||
|
}
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
var progressWriter progress.Writer
|
||||||
|
|
||||||
|
func StartProgressWriter() {
|
||||||
|
pw := progress.NewWriter()
|
||||||
|
|
||||||
|
pw.SetAutoStop(false)
|
||||||
|
pw.SetTrackerLength(25)
|
||||||
|
pw.SetMessageLength(20)
|
||||||
|
pw.SetSortBy(progress.SortByPercentDsc)
|
||||||
|
pw.SetStyle(progress.StyleDefault)
|
||||||
|
pw.SetOutputWriter(os.Stderr)
|
||||||
|
pw.SetTrackerPosition(progress.PositionRight)
|
||||||
|
pw.SetUpdateFrequency(time.Millisecond * 100)
|
||||||
|
pw.Style().Colors = progress.StyleColorsExample
|
||||||
|
pw.Style().Options.PercentFormat = "%4.1f%%"
|
||||||
|
pw.Style().Visibility.Pinned = true
|
||||||
|
pw.Style().Visibility.ETA = true
|
||||||
|
pw.Style().Visibility.Value = true
|
||||||
|
|
||||||
|
progressWriter = pw
|
||||||
|
go progressWriter.Render()
|
||||||
|
}
|
||||||
|
|
||||||
|
func StopProgressWriter() {
|
||||||
|
if progressWriter != nil {
|
||||||
|
progressWriter.Stop()
|
||||||
|
time.Sleep(1 * time.Second)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func SetPinnedMessageOnProgressWriter(msg string) {
|
||||||
|
if progressWriter != nil {
|
||||||
|
progressWriter.SetPinnedMessages(msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TrackProgress(message string, total int) ProgressTracker {
|
||||||
|
tracker := progress.Tracker{Message: message, Total: int64(total),
|
||||||
|
Units: progress.UnitsDefault}
|
||||||
|
|
||||||
|
if progressWriter != nil {
|
||||||
|
progressWriter.AppendTracker(&tracker)
|
||||||
|
}
|
||||||
|
|
||||||
|
return &progressTrackerImpl{tracker: &tracker}
|
||||||
|
}
|
||||||
|
|
||||||
|
func MarkTrackerAsDone(i any) {
|
||||||
|
if tracker, ok := i.(ProgressTracker); ok {
|
||||||
|
tracker.MarkAsDone()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func IncrementTrackerTotal(i any, count int64) {
|
||||||
|
if tracker, ok := i.(ProgressTracker); ok {
|
||||||
|
tracker.UpdateTotal(tracker.GetTotal() + count)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func IncrementProgress(i any, count int64) {
|
||||||
|
if tracker, ok := i.(ProgressTracker); ok && (progressTrackerDelta(tracker) > count) {
|
||||||
|
tracker.Increment(count)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func UpdateValue(i any, count int64) {
|
||||||
|
if tracker, ok := i.(ProgressTracker); ok {
|
||||||
|
tracker.SetValue(count)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func progressTrackerDelta(tracker ProgressTracker) int64 {
|
||||||
|
return (tracker.GetTotal() - tracker.GetValue())
|
||||||
|
}
|
||||||
@@ -9,14 +9,17 @@ import (
|
|||||||
"github.com/spf13/cobra"
|
"github.com/spf13/cobra"
|
||||||
)
|
)
|
||||||
|
|
||||||
func init() {
|
|
||||||
log.Init("pmg-logger", "debug")
|
|
||||||
}
|
|
||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
|
var debug bool
|
||||||
|
|
||||||
cmd := &cobra.Command{
|
cmd := &cobra.Command{
|
||||||
Use: "pmg",
|
Use: "pmg",
|
||||||
TraverseChildren: true,
|
TraverseChildren: true,
|
||||||
|
PersistentPreRun: func(cmd *cobra.Command, args []string) {
|
||||||
|
if debug {
|
||||||
|
log.Init("pmg-logger", "debug")
|
||||||
|
}
|
||||||
|
},
|
||||||
RunE: func(cmd *cobra.Command, args []string) error {
|
RunE: func(cmd *cobra.Command, args []string) error {
|
||||||
if len(args) == 0 {
|
if len(args) == 0 {
|
||||||
return nil
|
return nil
|
||||||
@@ -25,6 +28,8 @@ func main() {
|
|||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
|
cmd.PersistentFlags().BoolVar(&debug, "debug", false, "Enable debug logging")
|
||||||
|
|
||||||
cmd.AddCommand(ecosystems.NewNpmCommand())
|
cmd.AddCommand(ecosystems.NewNpmCommand())
|
||||||
|
|
||||||
if err := cmd.Execute(); err != nil {
|
if err := cmd.Execute(); err != nil {
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ import (
|
|||||||
packagev1 "buf.build/gen/go/safedep/api/protocolbuffers/go/safedep/messages/package/v1"
|
packagev1 "buf.build/gen/go/safedep/api/protocolbuffers/go/safedep/messages/package/v1"
|
||||||
malysisv1 "buf.build/gen/go/safedep/api/protocolbuffers/go/safedep/services/malysis/v1"
|
malysisv1 "buf.build/gen/go/safedep/api/protocolbuffers/go/safedep/services/malysis/v1"
|
||||||
"github.com/safedep/dry/log"
|
"github.com/safedep/dry/log"
|
||||||
|
"github.com/safedep/pmg/internal/ui"
|
||||||
"github.com/safedep/pmg/pkg/models"
|
"github.com/safedep/pmg/pkg/models"
|
||||||
vetUtils "github.com/safedep/vet/pkg/common/utils"
|
vetUtils "github.com/safedep/vet/pkg/common/utils"
|
||||||
)
|
)
|
||||||
@@ -19,6 +20,7 @@ type PackageAnalyser struct {
|
|||||||
Client malysisv1grpc.MalwareAnalysisServiceClient
|
Client malysisv1grpc.MalwareAnalysisServiceClient
|
||||||
Ctx context.Context
|
Ctx context.Context
|
||||||
MaliciousPkgsMutex sync.Mutex
|
MaliciousPkgsMutex sync.Mutex
|
||||||
|
ProgressTracker ui.ProgressTracker
|
||||||
}
|
}
|
||||||
|
|
||||||
func New(client malysisv1grpc.MalwareAnalysisServiceClient, ctx context.Context) *PackageAnalyser {
|
func New(client malysisv1grpc.MalwareAnalysisServiceClient, ctx context.Context) *PackageAnalyser {
|
||||||
@@ -66,6 +68,7 @@ func (ap *PackageAnalyser) Handler() vetUtils.WorkQueueFn[models.Package] {
|
|||||||
ap.MaliciousPkgsMutex.Unlock()
|
ap.MaliciousPkgsMutex.Unlock()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
ui.IncrementProgress(ap.ProgressTracker, 1)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user