Merge pull request #3589 from projectdiscovery/maint-headless-process

Removing redundant code with utils

v2/go.mod

@@ -32,7 +32,7 @@ require (
 	github.com/remeh/sizedwaitgroup v1.0.0
 	github.com/rs/xid v1.5.0
 	github.com/segmentio/ksuid v1.0.4
-	github.com/shirou/gopsutil/v3 v3.23.3
+	github.com/shirou/gopsutil/v3 v3.23.3 // indirect
 	github.com/spaolacci/murmur3 v1.1.0 // indirect
 	github.com/spf13/cast v1.5.0
 	github.com/syndtr/goleveldb v1.0.0
@@ -79,7 +79,7 @@ require (
 	github.com/projectdiscovery/sarif v0.0.1
 	github.com/projectdiscovery/tlsx v1.0.7
 	github.com/projectdiscovery/uncover v1.0.2
-	github.com/projectdiscovery/utils v0.0.21-0.20230419140949-a6527b072e4a
+	github.com/projectdiscovery/utils v0.0.24
 	github.com/projectdiscovery/wappalyzergo v0.0.88
 	github.com/stretchr/testify v1.8.2
 	gopkg.in/src-d/go-git.v4 v4.13.1

v2/go.sum

@@ -448,8 +448,8 @@ github.com/projectdiscovery/tlsx v1.0.7/go.mod h1:bMz1JMJf1sSBLpk7Y6vLwRZrijxPUv
 github.com/projectdiscovery/uncover v1.0.2 h1:mRFzflYyvwKkHd3XKufMlDRrb6p1mjFZTSHoNAUpFwo=
 github.com/projectdiscovery/uncover v1.0.2/go.mod h1:lz4QYfArSA6jJkXyB71kN2/Pc7IW7nJB8c95n7xtwqY=
 github.com/projectdiscovery/utils v0.0.3/go.mod h1:ne3eSlZlUKuhjHr8FfsfGcGteCzxcbJvFBx4VDBCxK0=
-github.com/projectdiscovery/utils v0.0.21-0.20230419140949-a6527b072e4a h1:h9ceITnnFLJ0qucXCrI3WOXtNRK2oAtcCuFsMAiuIc0=
+github.com/projectdiscovery/utils v0.0.24 h1:3kI8B5/ajKANPM7dfVHyTPFegy0OrVOxU+a3DXvX2a8=
-github.com/projectdiscovery/utils v0.0.21-0.20230419140949-a6527b072e4a/go.mod h1:954dxg9AWmNmcNQdc5BpucghibSvC76prWLQFrv14FQ=
+github.com/projectdiscovery/utils v0.0.24/go.mod h1:4ynwLqKugrMQzNjBJbzSDRBtadPwat/lwrXhWA6gdAE=
 github.com/projectdiscovery/wappalyzergo v0.0.88 h1:N/1vFlKmc3GJco9rANJdHrxg8jdav/xmnICo8rztmH8=
 github.com/projectdiscovery/wappalyzergo v0.0.88/go.mod h1:HvYuW0Be4JCjVds/+XAEaMSqRG9yrI97UmZq0TPk6A0=
 github.com/projectdiscovery/yamldoc-go v1.0.4 h1:eZoESapnMw6WAHiVgRwNqvbJEfNHEH148uthhFbG5jE=

v2/pkg/protocols/headless/engine/engine.go

@@ -9,13 +9,12 @@ import (
 	"github.com/go-rod/rod"
 	"github.com/go-rod/rod/lib/launcher"
 	"github.com/pkg/errors"
-	ps "github.com/shirou/gopsutil/v3/process"
 
 	"github.com/projectdiscovery/nuclei/v2/pkg/types"
 	fileutil "github.com/projectdiscovery/utils/file"
 	osutils "github.com/projectdiscovery/utils/os"
+	processutil "github.com/projectdiscovery/utils/process"
 	reflectutil "github.com/projectdiscovery/utils/reflect"
-	stringsutil "github.com/projectdiscovery/utils/strings"
 )
 
 // Browser is a browser structure for nuclei headless module
@@ -34,7 +33,7 @@ func New(options *types.Options) (*Browser, error) {
 	if err != nil {
 		return nil, errors.Wrap(err, "could not create temporary directory")
 	}
-	previousPIDs := findChromeProcesses()
+	previousPIDs := processutil.FindProcesses(processutil.IsChromeProcess)
 
 	chromeLauncher := launcher.New().
 		Leakless(false).
@@ -140,45 +139,5 @@ func (b *Browser) UserAgent() string {
 func (b *Browser) Close() {
 	b.engine.Close()
 	os.RemoveAll(b.tempDir)
-	b.killChromeProcesses()
+	processutil.CloseProcesses(processutil.IsChromeProcess, b.previousPIDs)
-}
-
-// killChromeProcesses any and all new chrome processes started after
-// headless process launch.
-func (b *Browser) killChromeProcesses() {
-	processes, _ := ps.Processes()
-
-	for _, process := range processes {
-		// skip non-chrome processes
-		if !isChromeProcess(process) {
-			continue
-		}
-		// skip chrome processes that were already running
-		if _, ok := b.previousPIDs[process.Pid]; ok {
-			continue
-		}
-		_ = process.Kill()
-	}
-}
-
-// findChromeProcesses finds chrome process running on host
-func findChromeProcesses() map[int32]struct{} {
-	processes, _ := ps.Processes()
-	list := make(map[int32]struct{})
-	for _, process := range processes {
-		if isChromeProcess(process) {
-			list[process.Pid] = struct{}{}
-			if ppid, err := process.Ppid(); err == nil {
-				list[ppid] = struct{}{}
-			}
-		}
-	}
-	return list
-}
-
-// isChromeProcess checks if a process is chrome/chromium
-func isChromeProcess(process *ps.Process) bool {
-	name, _ := process.Name()
-	executable, _ := process.Exe()
-	return stringsutil.ContainsAny(name, "chrome", "chromium") || stringsutil.ContainsAny(executable, "chrome", "chromium")
 }